diff --git a/.github/workflows/rust-ci.yml b/.github/workflows/rust-ci.yml index e249782..08f7c93 100644 --- a/.github/workflows/rust-ci.yml +++ b/.github/workflows/rust-ci.yml @@ -2,7 +2,7 @@ name: Rust on: push: - branches: [ "main" ] + branches: [ "main", "dev" ] pull_request: branches: [ "main" ] diff --git a/Cargo.toml b/Cargo.toml index ddc9467..2695885 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "usls" -version = "0.0.2" +version = "0.0.3" edition = "2021" description = "A Rust library integrated with ONNXRuntime, providing a collection of ML models." repository = "https://github.com/jamjamjon/usls" diff --git a/README.md b/README.md index 7668a97..219d5b3 100644 --- a/README.md +++ b/README.md @@ -1,18 +1,22 @@ # usls -A Rust library integrated with **ONNXRuntime**, providing a collection of **Computer Vison** and **Vision-Language** models including [YOLOv5](https://github.com/ultralytics/yolov5), [YOLOv8](https://github.com/ultralytics/ultralytics), [YOLOv9](https://github.com/WongKinYiu/yolov9), [RTDETR](https://arxiv.org/abs/2304.08069), [CLIP](https://github.com/openai/CLIP), [DINOv2](https://github.com/facebookresearch/dinov2), [FastSAM](https://github.com/CASIA-IVA-Lab/FastSAM), [YOLO-World](https://github.com/AILab-CVC/YOLO-World), [BLIP](https://arxiv.org/abs/2201.12086), [PaddleOCR](https://github.com/PaddlePaddle/PaddleOCR) , [Depth-Anything](https://github.com/LiheYoung/Depth-Anything) and others. +A Rust library integrated with **ONNXRuntime**, providing a collection of **Computer Vison** and **Vision-Language** models including [YOLOv5](https://github.com/ultralytics/yolov5), [YOLOv8](https://github.com/ultralytics/ultralytics), [YOLOv9](https://github.com/WongKinYiu/yolov9), [RTDETR](https://arxiv.org/abs/2304.08069), [CLIP](https://github.com/openai/CLIP), [DINOv2](https://github.com/facebookresearch/dinov2), [FastSAM](https://github.com/CASIA-IVA-Lab/FastSAM), [YOLO-World](https://github.com/AILab-CVC/YOLO-World), [BLIP](https://arxiv.org/abs/2201.12086), [PaddleOCR](https://github.com/PaddlePaddle/PaddleOCR), [Depth-Anything](https://github.com/LiheYoung/Depth-Anything), [MODNet](https://github.com/ZHKKKe/MODNet) and others. ## Recently Updated +| Portrait Matting | +| :----------------------------: | +|| + + | Depth-Anything | | :----------------------------: | || - | YOLOP-v2 | Face-Parsing | Text-Detection | | :----------------------------: | :------------------------------: | :------------------------------: | -|| | | +|| | | | YOLOv8-Obb | @@ -23,8 +27,6 @@ A Rust library integrated with **ONNXRuntime**, providing a collection of **Comp - - ## Supported Models | Model | Task / Type | Example | CUDA
f32 | CUDA
f16 | TensorRT
f32 | TensorRT
f16 | @@ -48,6 +50,7 @@ A Rust library integrated with **ONNXRuntime**, providing a collection of **Comp | [YOLOv5-classification](https://github.com/ultralytics/yolov5) | Object Detection | [demo](examples/yolov5) | ✅ | ✅ | ✅ | ✅ | | [YOLOv5-segmentation](https://github.com/ultralytics/yolov5) | Instance Segmentation | [demo](examples/yolov5) | ✅ | ✅ | ✅ | ✅ | | [Depth-Anything](https://github.com/LiheYoung/Depth-Anything) | Monocular Depth Estimation | [demo](examples/depth-anything) | ✅ | ✅ | ❌ | ❌ | +| [MODNet](https://github.com/ZHKKKe/MODNet) | Image Matting | [demo](examples/modnet) | ✅ | ✅ | ✅ | ✅ | ## Solution Models diff --git a/assets/portrait.jpg b/assets/portrait.jpg new file mode 100644 index 0000000..0ccd686 Binary files /dev/null and b/assets/portrait.jpg differ diff --git a/examples/blip/main.rs b/examples/blip/main.rs index b5b8439..51135f9 100644 --- a/examples/blip/main.rs +++ b/examples/blip/main.rs @@ -3,14 +3,14 @@ use usls::{models::Blip, DataLoader, Options}; fn main() -> Result<(), Box> { // visual let options_visual = Options::default() - .with_model("../models/blip-visual-base.onnx") + .with_model("blip-visual-base.onnx")? .with_i00((1, 1, 4).into()) .with_profile(false); // textual let options_textual = Options::default() - .with_model("../models/blip-textual-base.onnx") - .with_tokenizer("tokenizer-blip.json") + .with_model("blip-textual-base.onnx")? + .with_tokenizer("tokenizer-blip.json")? .with_i00((1, 1, 4).into()) // input_id: batch .with_i01((1, 1, 4).into()) // input_id: seq_len .with_i10((1, 1, 4).into()) // attention_mask: batch diff --git a/examples/clip/main.rs b/examples/clip/main.rs index a5e6a24..00f0bb0 100644 --- a/examples/clip/main.rs +++ b/examples/clip/main.rs @@ -3,14 +3,14 @@ use usls::{models::Clip, DataLoader, Options}; fn main() -> Result<(), Box> { // visual let options_visual = Options::default() - .with_model("../models/clip-b32-visual-dyn.onnx") + .with_model("clip-b32-visual-dyn.onnx")? .with_i00((1, 1, 4).into()) .with_profile(false); // textual let options_textual = Options::default() - .with_model("../models/clip-b32-textual-dyn.onnx") - .with_tokenizer("tokenizer-clip.json") + .with_model("clip-b32-textual-dyn.onnx")? + .with_tokenizer("tokenizer-clip.json")? .with_i00((1, 1, 4).into()) .with_profile(false); diff --git a/examples/db/main.rs b/examples/db/main.rs index f45a023..7b04f6d 100644 --- a/examples/db/main.rs +++ b/examples/db/main.rs @@ -10,7 +10,7 @@ fn main() -> Result<(), Box> { .with_min_width(5.0) .with_min_height(12.0) // .with_trt(0) - .with_model("../models/ppocr-v4-db-dyn.onnx"); + .with_model("ppocr-v4-db-dyn.onnx")?; let mut model = DB::new(&options)?; @@ -26,8 +26,8 @@ fn main() -> Result<(), Box> { // annotate let annotator = Annotator::default() .without_bboxes(true) - .with_masks_alpha(60) - .with_polygon_color([255, 105, 180, 255]) + .with_polygons_alpha(60) + .with_contours_color([255, 105, 180, 255]) .without_mbrs(true) .with_saveout("DB"); annotator.annotate(&x, &y); diff --git a/examples/depth-anything/demo.png b/examples/depth-anything/demo.png index 8325016..abf58e5 100644 Binary files a/examples/depth-anything/demo.png and b/examples/depth-anything/demo.png differ diff --git a/examples/depth-anything/main.rs b/examples/depth-anything/main.rs index 53dcca9..c4b70ae 100644 --- a/examples/depth-anything/main.rs +++ b/examples/depth-anything/main.rs @@ -3,7 +3,7 @@ use usls::{models::DepthAnything, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // visual let options = Options::default() - .with_model("../models/depth-anything-s-dyn.onnx") + .with_model("depth-anything-s-dyn.onnx")? .with_i00((1, 1, 8).into()) .with_i02((384, 512, 1024).into()) .with_i03((384, 512, 1024).into()); @@ -16,7 +16,9 @@ fn main() -> Result<(), Box> { let y = model.run(&x)?; // annotate - let annotator = Annotator::default().with_saveout("Depth-Anything"); + let annotator = Annotator::default() + .with_colormap("Turbo") + .with_saveout("Depth-Anything"); annotator.annotate(&x, &y); Ok(()) diff --git a/examples/dinov2/main.rs b/examples/dinov2/main.rs index 76310fa..5aa6682 100644 --- a/examples/dinov2/main.rs +++ b/examples/dinov2/main.rs @@ -3,7 +3,7 @@ use usls::{models::Dinov2, Options}; fn main() -> Result<(), Box> { // build model let options = Options::default() - .with_model("../models/dinov2-s14-dyn-f16.onnx") + .with_model("dinov2-s14-dyn-f16.onnx")? .with_i00((1, 1, 1).into()) .with_i02((224, 224, 224).into()) .with_i03((224, 224, 224).into()); diff --git a/examples/face-parsing/main.rs b/examples/face-parsing/main.rs index c2f20d8..47f4274 100644 --- a/examples/face-parsing/main.rs +++ b/examples/face-parsing/main.rs @@ -3,7 +3,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // build model let options = Options::default() - .with_model("../models/face-parsing-dyn.onnx") + .with_model("face-parsing-dyn.onnx")? .with_i00((1, 1, 4).into()) .with_i02((416, 640, 800).into()) .with_i03((416, 640, 800).into()) @@ -23,8 +23,8 @@ fn main() -> Result<(), Box> { .without_bboxes(true) .without_bboxes_conf(true) .without_bboxes_name(true) - .without_polygons(false) - .with_masks_name(false) + .without_contours(false) + .with_polygons_name(false) .with_saveout("Face-Parsing"); annotator.annotate(&x, &y); diff --git a/examples/fastsam/main.rs b/examples/fastsam/main.rs index 9359402..38af201 100644 --- a/examples/fastsam/main.rs +++ b/examples/fastsam/main.rs @@ -3,7 +3,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // build model let options = Options::default() - .with_model("../models/FastSAM-s-dyn-f16.onnx") + .with_model("FastSAM-s-dyn-f16.onnx")? .with_i00((1, 1, 4).into()) .with_i02((416, 640, 800).into()) .with_i03((416, 640, 800).into()) diff --git a/examples/modnet/README.md b/examples/modnet/README.md new file mode 100644 index 0000000..b866656 --- /dev/null +++ b/examples/modnet/README.md @@ -0,0 +1,15 @@ +## Quick Start + +```shell +cargo run -r --example modnet +``` + +## ONNX Model + +- [modnet-dyn](https://github.com/jamjamjon/assets/releases/download/v0.0.1/modnet-dyn.onnx) + + + +## Results + +![](./demo.png) diff --git a/examples/modnet/demo.png b/examples/modnet/demo.png new file mode 100644 index 0000000..9341e2f Binary files /dev/null and b/examples/modnet/demo.png differ diff --git a/examples/modnet/main.rs b/examples/modnet/main.rs new file mode 100644 index 0000000..0732064 --- /dev/null +++ b/examples/modnet/main.rs @@ -0,0 +1,23 @@ +use usls::{models::MODNet, Annotator, DataLoader, Options}; + +fn main() -> Result<(), Box> { + // build model + let options = Options::default() + .with_model("modnet-dyn.onnx")? + .with_i00((1, 1, 4).into()) + .with_i02((416, 512, 800).into()) + .with_i03((416, 512, 800).into()); + let model = MODNet::new(&options)?; + + // load image + let x = vec![DataLoader::try_read("./assets/portrait.jpg")?]; + + // run + let y = model.run(&x)?; + + // annotate + let annotator = Annotator::default().with_saveout("MODNet"); + annotator.annotate(&x, &y); + + Ok(()) +} diff --git a/examples/rtdetr/main.rs b/examples/rtdetr/main.rs index 4c6b413..4738604 100644 --- a/examples/rtdetr/main.rs +++ b/examples/rtdetr/main.rs @@ -3,7 +3,7 @@ use usls::{coco, models::RTDETR, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // build model let options = Options::default() - .with_model("../models/rtdetr-l-f16.onnx") + .with_model("rtdetr-l-f16.onnx")? .with_confs(&[0.4, 0.15]) // person: 0.4, others: 0.15 .with_names(&coco::NAMES_80); let mut model = RTDETR::new(&options)?; diff --git a/examples/rtmo/main.rs b/examples/rtmo/main.rs index 3aba47e..e43c6ec 100644 --- a/examples/rtmo/main.rs +++ b/examples/rtmo/main.rs @@ -3,7 +3,7 @@ use usls::{coco, models::RTMO, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // build model let options = Options::default() - .with_model("../rtmo-s-dyn.onnx") + .with_model("rtmo-s-dyn.onnx")? .with_i00((1, 1, 8).into()) .with_nk(17) .with_confs(&[0.3]) diff --git a/examples/svtr/main.rs b/examples/svtr/main.rs index e1a14dc..2eae102 100644 --- a/examples/svtr/main.rs +++ b/examples/svtr/main.rs @@ -6,8 +6,8 @@ fn main() -> Result<(), Box> { .with_i00((1, 2, 8).into()) .with_i03((320, 960, 1600).into()) .with_confs(&[0.2]) - .with_vocab("../ppocr_rec_vocab.txt") - .with_model("../models/ppocr-v4-svtr-ch-dyn.onnx"); + .with_vocab("ppocr_rec_vocab.txt")? + .with_model("ppocr-v4-svtr-ch-dyn.onnx")?; let mut model = SVTR::new(&options)?; // load images diff --git a/examples/yolo-world/main.rs b/examples/yolo-world/main.rs index 07de88a..f35cb49 100644 --- a/examples/yolo-world/main.rs +++ b/examples/yolo-world/main.rs @@ -3,7 +3,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // build model let options = Options::default() - .with_model("../models/yolov8s-world-v2-shoes.onnx") + .with_model("yolov8s-world-v2-shoes.onnx")? .with_i00((1, 1, 4).into()) .with_i02((416, 640, 800).into()) .with_i03((416, 640, 800).into()) diff --git a/examples/yolop/main.rs b/examples/yolop/main.rs index d7755fb..260d81f 100644 --- a/examples/yolop/main.rs +++ b/examples/yolop/main.rs @@ -3,7 +3,7 @@ use usls::{models::YOLOPv2, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // build model let options = Options::default() - .with_model("../models/yolopv2-dyn-480x800.onnx") + .with_model("yolopv2-dyn-480x800.onnx")? .with_i00((1, 1, 8).into()) .with_confs(&[0.3]); let mut model = YOLOPv2::new(&options)?; @@ -16,7 +16,7 @@ fn main() -> Result<(), Box> { // annotate let annotator = Annotator::default() - .with_masks_name(true) + .with_polygons_name(true) .with_saveout("YOLOPv2"); annotator.annotate(&x, &y); diff --git a/examples/yolov5/main.rs b/examples/yolov5/main.rs index 919fcea..4905bf4 100644 --- a/examples/yolov5/main.rs +++ b/examples/yolov5/main.rs @@ -9,7 +9,7 @@ fn main() -> Result<(), Box> { .with_conf_independent(true) .with_anchors_first(true) .with_yolo_task(YOLOTask::Segment) - .with_model("../models/yolov5s-seg.onnx") + .with_model("yolov5s-seg.onnx")? .with_trt(0) .with_fp16(true) .with_i00((1, 1, 4).into()) diff --git a/examples/yolov8-face/main.rs b/examples/yolov8-face/main.rs index 29984f2..c605869 100644 --- a/examples/yolov8-face/main.rs +++ b/examples/yolov8-face/main.rs @@ -3,7 +3,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // build model let options = Options::default() - .with_model("../models/yolov8n-face-dyn-f16.onnx") + .with_model("yolov8n-face-dyn-f16.onnx")? .with_i00((1, 1, 4).into()) .with_i02((416, 640, 800).into()) .with_i03((416, 640, 800).into()) diff --git a/examples/yolov8-falldown/main.rs b/examples/yolov8-falldown/main.rs index 9b44908..32824e2 100644 --- a/examples/yolov8-falldown/main.rs +++ b/examples/yolov8-falldown/main.rs @@ -2,7 +2,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // build model - let options = Options::default().with_model("../models/yolov8-falldown-f16.onnx"); + let options = Options::default().with_model("yolov8-falldown-f16.onnx")?; let mut model = YOLO::new(&options)?; // load image diff --git a/examples/yolov8-head/main.rs b/examples/yolov8-head/main.rs index b64531e..3280a68 100644 --- a/examples/yolov8-head/main.rs +++ b/examples/yolov8-head/main.rs @@ -2,7 +2,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // build model - let options = Options::default().with_model("../models/yolov8-head-f16.onnx"); + let options = Options::default().with_model("yolov8-head-f16.onnx")?; let mut model = YOLO::new(&options)?; // load image diff --git a/examples/yolov8-trash/main.rs b/examples/yolov8-trash/main.rs index 421ab44..91870db 100644 --- a/examples/yolov8-trash/main.rs +++ b/examples/yolov8-trash/main.rs @@ -3,7 +3,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // 1.build model let options = Options::default() - .with_model("../models/yolov8-plastic-bag-f16.onnx") + .with_model("yolov8-plastic-bag-f16.onnx")? .with_names(&["trash"]); let mut model = YOLO::new(&options)?; diff --git a/examples/yolov8/main.rs b/examples/yolov8/main.rs index 062eafd..f60cfb3 100644 --- a/examples/yolov8/main.rs +++ b/examples/yolov8/main.rs @@ -3,23 +3,20 @@ use usls::{coco, models::YOLO, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // build model let options = Options::default() - // .with_model("../models/yolov8m.onnx") - // .with_model("../models/yolov8m-dyn-f16.onnx") - // .with_model("../models/yolov8m-pose-dyn-f16.onnx") - // .with_model("../models/yolov8m-seg-dyn-f16.onnx") - .with_model("../models/yolov8s-cls.onnx") - // .with_model("../models/yolov8s-obb.onnx") + .with_model("yolov8m-dyn.onnx")? + // .with_model("yolov8m-pose-dyn.onnx")? + // .with_model("yolov8m-cls-dyn.onnx")? + // .with_model("yolov8m-seg-dyn.onnx")? + // .with_model("yolov8m-obb-dyn.onnx")? + // .with_model("yolov8m-oiv7-dyn.onnx")? // .with_trt(0) // .with_fp16(true) .with_i00((1, 1, 4).into()) - .with_i02((224, 1024, 1024).into()) - .with_i03((224, 1024, 1024).into()) - // .with_i02((224, 640, 800).into()) - // .with_i03((224, 640, 800).into()) + .with_i02((224, 640, 800).into()) + .with_i03((224, 640, 800).into()) .with_confs(&[0.4, 0.15]) // person: 0.4, others: 0.15 .with_names2(&coco::KEYPOINTS_NAMES_17) - .with_profile(true) - .with_dry_run(10); + .with_profile(false); let mut model = YOLO::new(&options)?; // build dataloader @@ -30,41 +27,7 @@ fn main() -> Result<(), Box> { // build annotate let annotator = Annotator::default() - // .with_probs_topk(10) - // // bboxes - // .without_bboxes(false) - // .without_bboxes_conf(false) - // .without_bboxes_name(false) - // .without_bboxes_text_bg(false) - // .with_bboxes_text_color([255, 255, 255, 255]) - // .with_bboxes_text_bg_alpha(255) - // // keypoints - // .without_keypoints(false) - // .with_keypoints_palette(&COCO_KEYPOINT_COLORS_17) .with_skeletons(&coco::SKELETONS_16) - // .with_keypoints_name(false) - // .with_keypoints_conf(false) - // .without_keypoints_text_bg(false) - // .with_keypoints_text_color([255, 255, 255, 255]) - // .with_keypoints_text_bg_alpha(255) - // .with_keypoints_radius(4) - // // masks - // .without_masks(false) - // .with_masks_alpha(190) - // .without_polygons(false) - // // .with_polygon_color([0, 255, 255, 255]) - // .with_masks_conf(false) - // .with_masks_name(true) - // .with_masks_text_bg(true) - // .with_masks_text_color([255, 255, 255, 255]) - // .with_masks_text_bg_alpha(10) - // // mbrs - // .without_mbrs(false) - // .without_mbrs_conf(false) - // .without_mbrs_name(false) - // .without_mbrs_text_bg(false) - // .with_mbrs_text_color([255, 255, 255, 255]) - // .with_mbrs_text_bg_alpha(70) .with_saveout("YOLOv8"); // run & annotate diff --git a/examples/yolov9/main.rs b/examples/yolov9/main.rs index fd1e528..6f6d8bf 100644 --- a/examples/yolov9/main.rs +++ b/examples/yolov9/main.rs @@ -3,7 +3,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // build model let options = Options::default() - .with_model("../models/yolov9-c-dyn-f16.onnx") + .with_model("yolov9-c-dyn-f16.onnx")? .with_i00((1, 1, 4).into()) .with_i02((416, 640, 800).into()) .with_i03((416, 640, 800).into()) diff --git a/src/core/annotator.rs b/src/core/annotator.rs index 3d4fa3a..e71a7b9 100644 --- a/src/core/annotator.rs +++ b/src/core/annotator.rs @@ -1,9 +1,10 @@ use crate::{ - auto_load, string_now, Bbox, Keypoint, Mask, Mbr, Prob, CHECK_MARK, CROSS_MARK, TURBO, Y, + auto_load, colormap256, string_now, Bbox, Keypoint, Mask, Mbr, Polygon, Prob, CHECK_MARK, + CROSS_MARK, Y, }; use ab_glyph::{FontVec, PxScale}; use anyhow::Result; -use image::{DynamicImage, ImageBuffer, Rgba, RgbaImage}; +use image::{DynamicImage, GenericImage, Rgba, RgbaImage}; use imageproc::map::map_colors; /// Annotator for struct `Y` @@ -13,18 +14,21 @@ pub struct Annotator { _scale: f32, // Cope with ab_glyph & imageproc=0.24.0 scale_dy: f32, saveout: Option, + // About mbrs without_mbrs: bool, without_mbrs_conf: bool, without_mbrs_name: bool, without_mbrs_text_bg: bool, mbrs_text_color: Rgba, + // About bboxes without_bboxes: bool, without_bboxes_conf: bool, without_bboxes_name: bool, without_bboxes_text_bg: bool, bboxes_text_color: Rgba, + // About keypoints without_keypoints: bool, with_keypoints_conf: bool, @@ -34,15 +38,21 @@ pub struct Annotator { skeletons: Option>, keypoints_radius: usize, keypoints_palette: Option>, + + // About polygons + without_polygons: bool, + without_contours: bool, + with_polygons_conf: bool, + with_polygons_name: bool, + with_polygons_text_bg: bool, + polygons_text_color: Rgba, + polygons_alpha: u8, + contours_color: Rgba, + // About masks without_masks: bool, - without_polygons: bool, - with_masks_conf: bool, - with_masks_name: bool, - with_masks_text_bg: bool, - masks_text_color: Rgba, - masks_alpha: u8, - polygon_color: Rgba, + colormap: Option<[[u8; 3]; 256]>, + // About probs probs_topk: usize, } @@ -53,7 +63,7 @@ impl Default for Annotator { font: Self::load_font(None).unwrap(), _scale: 6.666667, scale_dy: 28., - masks_alpha: 179, + polygons_alpha: 179, saveout: None, without_bboxes: false, without_bboxes_conf: false, @@ -73,19 +83,22 @@ impl Default for Annotator { keypoints_palette: None, without_keypoints_text_bg: false, keypoints_text_color: Rgba([0, 0, 0, 255]), - without_masks: false, without_polygons: false, - polygon_color: Rgba([255, 255, 255, 255]), - with_masks_name: false, - with_masks_conf: false, - with_masks_text_bg: false, - masks_text_color: Rgba([255, 255, 255, 255]), + without_contours: false, + contours_color: Rgba([255, 255, 255, 255]), + with_polygons_name: false, + with_polygons_conf: false, + with_polygons_text_bg: false, + polygons_text_color: Rgba([255, 255, 255, 255]), probs_topk: 5usize, + without_masks: false, + colormap: None, } } } impl Annotator { + /// Plotting BBOXes or not pub fn without_bboxes(mut self, x: bool) -> Self { self.without_bboxes = x; self @@ -161,6 +174,7 @@ impl Annotator { self } + /// Plotting MBRs or not pub fn without_mbrs(mut self, x: bool) -> Self { self.without_mbrs = x; self @@ -191,48 +205,68 @@ impl Annotator { self } - pub fn without_masks(mut self, x: bool) -> Self { - self.without_masks = x; - self - } - pub fn without_polygons(mut self, x: bool) -> Self { self.without_polygons = x; self } - pub fn with_masks_conf(mut self, x: bool) -> Self { - self.with_masks_conf = x; + pub fn without_contours(mut self, x: bool) -> Self { + self.without_contours = x; self } - pub fn with_masks_name(mut self, x: bool) -> Self { - self.with_masks_name = x; + pub fn with_polygons_conf(mut self, x: bool) -> Self { + self.with_polygons_conf = x; self } - pub fn with_masks_text_bg(mut self, x: bool) -> Self { - self.with_masks_text_bg = x; + pub fn with_polygons_name(mut self, x: bool) -> Self { + self.with_polygons_name = x; self } - pub fn with_masks_text_color(mut self, rgba: [u8; 4]) -> Self { - self.masks_text_color = Rgba(rgba); + pub fn with_polygons_text_bg(mut self, x: bool) -> Self { + self.with_polygons_text_bg = x; self } - pub fn with_masks_alpha(mut self, x: u8) -> Self { - self.masks_alpha = x; + pub fn with_colormap(mut self, x: &str) -> Self { + let x = match x { + "turbo" | "Turbo" | "TURBO" => colormap256::TURBO, + "inferno" | "Inferno" | "INFERNO" => colormap256::INFERNO, + "plasma" | "Plasma" | "PLASMA" => colormap256::PLASMA, + "viridis" | "Viridis" | "VIRIDIS" => colormap256::VIRIDIS, + "magma" | "Magma" | "MAGMA" => colormap256::MAGMA, + "bentcoolwarm" | "BentCoolWarm" | "BENTCOOLWARM" => colormap256::BENTCOOLWARM, + "blackbody" | "BlackBody" | "BLACKBODY" => colormap256::BLACKBODY, + "extendedkindLmann" | "ExtendedKindLmann" | "EXTENDEDKINDLMANN" => { + colormap256::EXTENDEDKINDLMANN + } + "kindlmann" | "KindLmann" | "KINDLMANN" => colormap256::KINDLMANN, + "smoothcoolwarm" | "SmoothCoolWarm" | "SMOOTHCOOLWARM" => colormap256::SMOOTHCOOLWARM, + _ => todo!(), + }; + self.colormap = Some(x); self } - pub fn with_masks_text_bg_alpha(mut self, x: u8) -> Self { - self.masks_text_color.0[3] = x; + pub fn with_polygons_text_color(mut self, rgba: [u8; 4]) -> Self { + self.polygons_text_color = Rgba(rgba); self } - pub fn with_polygon_color(mut self, rgba: [u8; 4]) -> Self { - self.polygon_color = Rgba(rgba); + pub fn with_polygons_alpha(mut self, x: u8) -> Self { + self.polygons_alpha = x; + self + } + + pub fn with_polygons_text_bg_alpha(mut self, x: u8) -> Self { + self.polygons_text_color.0[3] = x; + self + } + + pub fn with_contours_color(mut self, rgba: [u8; 4]) -> Self { + self.contours_color = Rgba(rgba); self } @@ -251,6 +285,7 @@ impl Annotator { self } + /// Save annotated images to `runs` folder pub fn save(&self, image: &RgbaImage, saveout: &str) { let mut saveout = std::path::PathBuf::from("runs").join(saveout); if !saveout.exists() { @@ -264,60 +299,62 @@ impl Annotator { } } + /// Annotate images pub fn annotate(&self, imgs: &[DynamicImage], ys: &[Y]) { for (img, y) in imgs.iter().zip(ys.iter()) { - let mut img_rgb = img.to_rgba8(); + let mut img_rgba = img.to_rgba8(); - // pixels - if !self.without_masks { - if let Some(xs) = &y.pixels() { - self.plot_pixels(&mut img_rgb, xs) - } - } - - // masks - if !self.without_masks { - if let Some(xs) = &y.masks() { - self.plot_masks_and_polygons(&mut img_rgb, xs) + // polygons + if !self.without_polygons { + if let Some(xs) = &y.polygons() { + self.plot_polygons(&mut img_rgba, xs) } } // bboxes if !self.without_bboxes { if let Some(xs) = &y.bboxes() { - self.plot_bboxes(&mut img_rgb, xs) + self.plot_bboxes(&mut img_rgba, xs) } } // mbrs if !self.without_mbrs { if let Some(xs) = &y.mbrs() { - self.plot_mbrs(&mut img_rgb, xs) + self.plot_mbrs(&mut img_rgba, xs) } } // keypoints if !self.without_keypoints { if let Some(xs) = &y.keypoints() { - self.plot_keypoints(&mut img_rgb, xs) + self.plot_keypoints(&mut img_rgba, xs) } } // probs if let Some(xs) = &y.probs() { - self.plot_probs(&mut img_rgb, xs) + self.plot_probs(&mut img_rgba, xs) + } + + // masks + if !self.without_masks { + if let Some(xs) = &y.masks() { + self.plot_masks(&mut img_rgba, xs) + } } // save if let Some(saveout) = &self.saveout { - self.save(&img_rgb, saveout); + self.save(&img_rgba, saveout); } } } + /// Plot bounding bboxes and labels pub fn plot_bboxes(&self, img: &mut RgbaImage, bboxes: &[Bbox]) { for bbox in bboxes.iter() { - // bboxes + // bbox imageproc::drawing::draw_hollow_rect_mut( img, imageproc::rect::Rect::at(bbox.xmin().round() as i32, bbox.ymin().round() as i32) @@ -325,33 +362,26 @@ impl Annotator { image::Rgba(self.get_color(bbox.id() as usize).into()), ); - // texts - let mut legend = String::new(); - if !self.without_bboxes_name { - legend.push_str(&bbox.name().unwrap_or(&bbox.id().to_string()).to_string()); + // label + if !self.without_bboxes_name || !self.without_bboxes_conf { + let label = bbox.label(!self.without_bboxes_name, !self.without_bboxes_conf); + self.put_text( + img, + &label, + bbox.xmin(), + bbox.ymin(), + image::Rgba(self.get_color(bbox.id() as usize).into()), + self.bboxes_text_color, + self.without_bboxes_text_bg, + ); } - if !self.without_bboxes_conf { - if !self.without_bboxes_name { - legend.push_str(&format!(": {:.4}", bbox.confidence())); - } else { - legend.push_str(&format!("{:.4}", bbox.confidence())); - } - } - self.put_text( - img, - legend.as_str(), - bbox.xmin(), - bbox.ymin(), - image::Rgba(self.get_color(bbox.id() as usize).into()), - self.bboxes_text_color, - self.without_bboxes_text_bg, - ); } } + /// Plot minimum bounding rectangle and labels pub fn plot_mbrs(&self, img: &mut RgbaImage, mbrs: &[Mbr]) { for mbr in mbrs.iter() { - // mbrs + // mbr for i in 0..mbr.vertices().len() { let p1 = mbr.vertices()[i]; let p2 = mbr.vertices()[(i + 1) % mbr.vertices().len()]; @@ -363,152 +393,204 @@ impl Annotator { ); } - // text - let mut legend = String::new(); - if !self.without_mbrs_name { - legend.push_str(&mbr.name().unwrap_or(&mbr.id().to_string()).to_string()); - } - if !self.without_mbrs_conf { - if !self.without_mbrs_name { - legend.push_str(&format!(": {:.4}", mbr.confidence())); - } else { - legend.push_str(&format!("{:.4}", mbr.confidence())); - } - } - self.put_text( - img, - legend.as_str(), - mbr.top().x as f32, - mbr.top().y as f32, - image::Rgba(self.get_color(mbr.id() as usize).into()), - self.mbrs_text_color, - self.without_mbrs_text_bg, - ); - } - } - - pub fn plot_pixels(&self, img: &mut RgbaImage, pixels: &[u8]) { - let (w, h) = img.dimensions(); - let luma: ImageBuffer, Vec> = - ImageBuffer::from_raw(w, h, pixels.to_vec()) - .expect("Faild to create luma from ndarray"); - let luma = map_colors(&luma, |p| { - let x = p[0]; - image::Rgb(TURBO[x as usize]) - }); - let luma = image::DynamicImage::from(luma); - let luma = luma.resize_exact(w / 2, h / 2, image::imageops::FilterType::CatmullRom); - let im_ori = img.clone(); - let im_ori = image::DynamicImage::from(im_ori); - let im_ori = im_ori.resize_exact(w / 2, h / 2, image::imageops::FilterType::CatmullRom); - - // overwrite - for x in 0..w { - for y in 0..h { - img.put_pixel(x, y, Rgba([255, 255, 255, 255])); - } - } - - // paste - let pos_x = 0; - let pos_y = (2 * (h - im_ori.height()) / 3) as i64; - image::imageops::overlay(img, &im_ori, pos_x, pos_y); - image::imageops::overlay(img, &luma, im_ori.width().into(), pos_y); - - // text - let legend = "Raw"; - let scale = PxScale::from(self.scale_dy * 2.5); - let (text_w, text_h) = imageproc::drawing::text_size(scale, &self.font, legend); - imageproc::drawing::draw_text_mut( - img, - Rgba([0, 0, 0, 255]), - ((im_ori.width() - text_w) / 2) as i32, - ((pos_y as u32 - text_h) / 2) as i32, - scale, - &self.font, - legend, - ); - let legend = "Depth"; - let (text_w, text_h) = imageproc::drawing::text_size(scale, &self.font, legend); - imageproc::drawing::draw_text_mut( - img, - Rgba([0, 0, 0, 255]), - (im_ori.width() + (im_ori.width() - text_w) / 2) as i32, - ((pos_y as u32 - text_h) / 2) as i32, - scale, - &self.font, - legend, - ); - } - - pub fn plot_masks_and_polygons(&self, img: &mut RgbaImage, masks: &[Mask]) { - let mut convas = img.clone(); - for mask in masks.iter() { - // masks - let polygon_i32 = mask - .polygon() - .exterior() - .points() - .take(if mask.is_closed() { - mask.count() - 1 - } else { - mask.count() - }) - .map(|p| imageproc::point::Point::new(p.x() as i32, p.y() as i32)) - .collect::>(); - let mut mask_color = self.get_color(mask.id() as usize); - mask_color.3 = self.masks_alpha; - imageproc::drawing::draw_polygon_mut( - &mut convas, - &polygon_i32, - Rgba(mask_color.into()), - ); - - // contours(polygons) - if !self.without_polygons { - let polygon_f32 = mask - .polygon() - .exterior() - .points() - .take(if mask.is_closed() { - mask.count() - 1 - } else { - mask.count() - }) - .map(|p| imageproc::point::Point::new(p.x() as f32, p.y() as f32)) - .collect::>(); - imageproc::drawing::draw_hollow_polygon_mut(img, &polygon_f32, self.polygon_color); - } - } - image::imageops::overlay(img, &convas, 0, 0); - - // text on top - for mask in masks.iter() { - if let Some((x, y)) = mask.centroid() { - let mut legend = String::new(); - if self.with_masks_name { - legend.push_str(&mask.name().unwrap_or(&mask.id().to_string()).to_string()); - } - if self.with_masks_conf { - if self.with_masks_name { - legend.push_str(&format!(": {:.4}", mask.confidence())); - } else { - legend.push_str(&format!("{:.4}", mask.confidence())); - } - } + // label + if !self.without_mbrs_name || !self.without_mbrs_conf { + let label = mbr.label(!self.without_mbrs_name, !self.without_mbrs_conf); self.put_text( img, - legend.as_str(), - x, - y, - image::Rgba(self.get_color(mask.id() as usize).into()), - self.masks_text_color, - !self.with_masks_text_bg, + &label, + mbr.top().x as f32, + mbr.top().y as f32, + image::Rgba(self.get_color(mbr.id() as usize).into()), + self.mbrs_text_color, + self.without_mbrs_text_bg, ); } } } + /// Plot polygons(hollow & filled) and labels + pub fn plot_polygons(&self, img: &mut RgbaImage, polygons: &[Polygon]) { + let mut convas = img.clone(); + for polygon in polygons.iter() { + // filled + let polygon_i32 = polygon + .polygon() + .exterior() + .points() + .take(if polygon.is_closed() { + polygon.count() - 1 + } else { + polygon.count() + }) + .map(|p| imageproc::point::Point::new(p.x() as i32, p.y() as i32)) + .collect::>(); + let mut color_ = self.get_color(polygon.id() as usize); + color_.3 = self.polygons_alpha; + imageproc::drawing::draw_polygon_mut(&mut convas, &polygon_i32, Rgba(color_.into())); + + // contour + if !self.without_contours { + let polygon_f32 = polygon + .polygon() + .exterior() + .points() + .take(if polygon.is_closed() { + polygon.count() - 1 + } else { + polygon.count() + }) + .map(|p| imageproc::point::Point::new(p.x() as f32, p.y() as f32)) + .collect::>(); + imageproc::drawing::draw_hollow_polygon_mut(img, &polygon_f32, self.contours_color); + } + } + image::imageops::overlay(img, &convas, 0, 0); + + // labels on top + if self.with_polygons_name || self.with_polygons_conf { + for polygon in polygons.iter() { + if let Some((x, y)) = polygon.centroid() { + let label = polygon.label(self.with_polygons_name, self.with_polygons_conf); + self.put_text( + img, + &label, + x, + y, + image::Rgba(self.get_color(polygon.id() as usize).into()), + self.polygons_text_color, + !self.with_polygons_text_bg, + ); + } + } + } + } + + /// Plot keypoints and texts + pub fn plot_keypoints(&self, img: &mut RgbaImage, keypoints: &[Vec]) { + for kpts in keypoints.iter() { + for (i, kpt) in kpts.iter().enumerate() { + if kpt.confidence() == 0.0 { + continue; + } + + // keypoint + let color = match &self.keypoints_palette { + None => self.get_color(i + 10), + Some(keypoints_palette) => keypoints_palette[i], + }; + imageproc::drawing::draw_filled_circle_mut( + img, + (kpt.x() as i32, kpt.y() as i32), + self.keypoints_radius as i32, + image::Rgba(color.into()), + ); + + // label + if self.with_keypoints_name || self.with_keypoints_conf { + let label = kpt.label(self.with_keypoints_name, self.with_keypoints_conf); + self.put_text( + img, + &label, + kpt.x(), + kpt.y(), + image::Rgba(self.get_color(kpt.id() as usize).into()), + self.keypoints_text_color, + self.without_keypoints_text_bg, + ); + } + } + + // skeletons + if let Some(skeletons) = &self.skeletons { + for &(i, ii) in skeletons.iter() { + let kpt1 = &kpts[i]; + let kpt2 = &kpts[ii]; + if kpt1.confidence() == 0.0 || kpt2.confidence() == 0.0 { + continue; + } + imageproc::drawing::draw_line_segment_mut( + img, + (kpt1.x(), kpt1.y()), + (kpt2.x(), kpt2.y()), + image::Rgba([255, 51, 255, 255]), + ); + } + } + } + } + + /// Plot masks + pub fn plot_masks(&self, img: &mut RgbaImage, masks: &[Mask]) { + let (w, h) = img.dimensions(); + // let hstack = w < h; + let hstack = true; + let scale = 2; + let size = (masks.len() + 1) as u32; + + // convas + let convas = img.clone(); + let mut convas = image::DynamicImage::from(convas); + if hstack { + convas = convas.resize_exact( + w, + h / scale * (size / scale), + image::imageops::FilterType::CatmullRom, + ); + } else { + convas = convas.resize_exact( + w / scale, + h * size / scale, + image::imageops::FilterType::CatmullRom, + ); + } + for x in 0..convas.width() { + for y in 0..convas.height() { + convas.put_pixel(x, y, Rgba([255, 255, 255, 255])); + } + } + + // place original + let im_ori = img.clone(); + let im_ori = image::DynamicImage::from(im_ori); + let im_ori = im_ori.resize_exact( + w / scale, + h / scale, + image::imageops::FilterType::CatmullRom, + ); + image::imageops::overlay(&mut convas, &im_ori, 0, 0); + + // place masks + for (i, mask) in masks.iter().enumerate() { + let i = i + 1; + let luma = if let Some(colormap) = self.colormap { + let luma = map_colors(mask.mask(), |p| { + let x = p[0]; + image::Rgb(colormap[x as usize]) + }); + image::DynamicImage::from(luma) + } else { + mask.mask().to_owned() + }; + let luma = luma.resize_exact( + w / scale, + h / scale, + image::imageops::FilterType::CatmullRom, + ); + if hstack { + let pos_x = (i as u32 % scale) * luma.width(); + let pos_y = (i as u32 / scale) * luma.height(); + image::imageops::overlay(&mut convas, &luma, pos_x as i64, pos_y as i64); + } else { + let pos_x = 0; + let pos_y = i as u32 * luma.height(); + image::imageops::overlay(&mut convas, &luma, pos_x as i64, pos_y as i64); + } + } + *img = convas.into_rgba8(); + } + + /// Plot probs pub fn plot_probs(&self, img: &mut RgbaImage, probs: &Prob) { let (x, mut y) = (img.width() as i32 / 20, img.height() as i32 / 20); for k in probs.topk(self.probs_topk).iter() { @@ -534,67 +616,7 @@ impl Annotator { } } - pub fn plot_keypoints(&self, img: &mut RgbaImage, keypoints: &[Vec]) { - for kpts in keypoints.iter() { - for (i, kpt) in kpts.iter().enumerate() { - if kpt.confidence() == 0.0 { - continue; - } - - // keypoints - let color = match &self.keypoints_palette { - None => self.get_color(i + 10), - Some(keypoints_palette) => keypoints_palette[i], - }; - imageproc::drawing::draw_filled_circle_mut( - img, - (kpt.x() as i32, kpt.y() as i32), - self.keypoints_radius as i32, - image::Rgba(color.into()), - ); - - // text - let mut legend = String::new(); - if self.with_keypoints_name { - legend.push_str(&kpt.name().unwrap_or(&kpt.id().to_string()).to_string()); - } - if self.with_keypoints_conf { - if self.with_keypoints_name { - legend.push_str(&format!(": {:.4}", kpt.confidence())); - } else { - legend.push_str(&format!("{:.4}", kpt.confidence())); - } - } - self.put_text( - img, - legend.as_str(), - kpt.x(), - kpt.y(), - image::Rgba(self.get_color(kpt.id() as usize).into()), - self.keypoints_text_color, - self.without_keypoints_text_bg, - ); - } - - // draw skeleton - if let Some(skeletons) = &self.skeletons { - for &(i, ii) in skeletons.iter() { - let kpt1 = &kpts[i]; - let kpt2 = &kpts[ii]; - if kpt1.confidence() == 0.0 || kpt2.confidence() == 0.0 { - continue; - } - imageproc::drawing::draw_line_segment_mut( - img, - (kpt1.x(), kpt1.y()), - (kpt2.x(), kpt2.y()), - image::Rgba([255, 51, 255, 255]), - ); - } - } - } - } - + /// Helper for putting texts #[allow(clippy::too_many_arguments)] fn put_text( &self, @@ -642,19 +664,22 @@ impl Annotator { } } + /// Load custom font fn load_font(path: Option<&str>) -> Result { let path_font = match path { - None => auto_load("Arial.ttf")?, + None => auto_load("Arial.ttf", Some("fonts"))?, Some(p) => p.into(), }; let buffer = std::fs::read(path_font)?; Ok(FontVec::try_from_vec(buffer.to_owned()).unwrap()) } + /// Pick color from pallette pub fn get_color(&self, n: usize) -> (u8, u8, u8, u8) { Self::color_palette()[n % Self::color_palette().len()] } + /// Color pallette fn color_palette() -> [(u8, u8, u8, u8); 20] { [ (0, 255, 127, 255), // spring green diff --git a/src/core/engine.rs b/src/core/engine.rs index 58412d2..eba7bd6 100644 --- a/src/core/engine.rs +++ b/src/core/engine.rs @@ -6,7 +6,7 @@ use ort::{ TensorRTExecutionProvider, ValueType, }; -use crate::{config_dir, Device, MinOptMax, Options, CHECK_MARK, CROSS_MARK, SAFE_CROSS_MARK}; +use crate::{home_dir, Device, MinOptMax, Options, CHECK_MARK, CROSS_MARK, SAFE_CROSS_MARK}; /// ONNXRuntime Backend #[derive(Debug)] @@ -211,7 +211,7 @@ impl OrtEngine { .with_engine_cache(engine_cache_enable) .with_engine_cache_path(format!( "{}/{}", - config_dir().to_str().unwrap(), + home_dir(None).to_str().unwrap(), "trt-cache" )) .with_timing_cache(false) diff --git a/src/core/options.rs b/src/core/options.rs index 12d1d66..f6d6168 100644 --- a/src/core/options.rs +++ b/src/core/options.rs @@ -1,3 +1,5 @@ +use anyhow::Result; + use crate::{auto_load, models::YOLOTask, Device, MinOptMax}; /// Options for building models @@ -117,9 +119,9 @@ impl Default for Options { } impl Options { - pub fn with_model(mut self, onnx_path: &str) -> Self { - self.onnx_path = auto_load(onnx_path).unwrap(); - self + pub fn with_model(mut self, onnx_path: &str) -> Result { + self.onnx_path = auto_load(onnx_path, Some("models"))?; + Ok(self) } pub fn with_dry_run(mut self, n: usize) -> Self { @@ -187,14 +189,14 @@ impl Options { self } - pub fn with_vocab(mut self, vocab: &str) -> Self { - self.vocab = Some(auto_load(vocab).unwrap()); - self + pub fn with_vocab(mut self, vocab: &str) -> Result { + self.vocab = Some(auto_load(vocab, Some("models"))?); + Ok(self) } - pub fn with_tokenizer(mut self, tokenizer: &str) -> Self { - self.tokenizer = Some(auto_load(tokenizer).unwrap()); - self + pub fn with_tokenizer(mut self, tokenizer: &str) -> Result { + self.tokenizer = Some(auto_load(tokenizer, Some("models"))?); + Ok(self) } pub fn with_unclip_ratio(mut self, x: f32) -> Self { diff --git a/src/models/db.rs b/src/models/db.rs index 84ecfd9..76dd04a 100644 --- a/src/models/db.rs +++ b/src/models/db.rs @@ -1,4 +1,4 @@ -use crate::{ops, DynConf, Mask, Mbr, MinOptMax, Options, OrtEngine, Y}; +use crate::{ops, DynConf, Mbr, MinOptMax, Options, OrtEngine, Polygon, Y}; use anyhow::Result; use image::DynamicImage; use ndarray::{Array, Axis, IxDyn}; @@ -56,7 +56,7 @@ impl DB { let mut ys = Vec::new(); for (idx, luma) in xs[0].axis_iter(Axis(0)).enumerate() { let mut y_bbox = Vec::new(); - let mut y_masks: Vec = Vec::new(); + let mut y_polygons: Vec = Vec::new(); let mut y_mbrs: Vec = Vec::new(); // reshape @@ -99,7 +99,7 @@ impl DB { { continue; } - let mask = Mask::default().with_points_imageproc(&contour.points); + let mask = Polygon::default().with_points_imageproc(&contour.points); let delta = mask.area() * ratio.round() as f64 * self.unclip_ratio as f64 / mask.perimeter(); let mask = mask @@ -120,7 +120,7 @@ impl DB { if let Some(mbr) = mask.mbr() { y_mbrs.push(mbr.with_confidence(confidence).with_id(0)); } - y_masks.push(mask.with_id(0)); + y_polygons.push(mask.with_id(0)); } else { continue; } @@ -128,7 +128,7 @@ impl DB { ys.push( Y::default() .with_bboxes(&y_bbox) - .with_masks(&y_masks) + .with_polygons(&y_polygons) .with_mbrs(&y_mbrs), ); } diff --git a/src/models/depth_anything.rs b/src/models/depth_anything.rs index ee3aa8f..f8dc406 100644 --- a/src/models/depth_anything.rs +++ b/src/models/depth_anything.rs @@ -1,4 +1,4 @@ -use crate::{ops, MinOptMax, Options, OrtEngine, Y}; +use crate::{ops, Mask, MinOptMax, Options, OrtEngine, Y}; use anyhow::Result; use image::{DynamicImage, ImageBuffer}; use ndarray::{Array, Axis, IxDyn}; @@ -59,7 +59,7 @@ impl DepthAnything { xs0[idx].height(), image::imageops::FilterType::CatmullRom, ); - ys.push(Y::default().with_pixels(&luma.into_luma8().into_raw())); + ys.push(Y::default().with_masks(&[Mask::default().with_mask(luma)])); } Ok(ys) } diff --git a/src/models/mod.rs b/src/models/mod.rs index 98e1e31..5b521f6 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -3,6 +3,7 @@ mod clip; mod db; mod depth_anything; mod dinov2; +mod modnet; mod rtdetr; mod rtmo; mod svtr; @@ -14,6 +15,7 @@ pub use clip::Clip; pub use db::DB; pub use depth_anything::DepthAnything; pub use dinov2::Dinov2; +pub use modnet::MODNet; pub use rtdetr::RTDETR; pub use rtmo::RTMO; pub use svtr::SVTR; diff --git a/src/models/modnet.rs b/src/models/modnet.rs new file mode 100644 index 0000000..5624b7b --- /dev/null +++ b/src/models/modnet.rs @@ -0,0 +1,76 @@ +use anyhow::Result; +use image::DynamicImage; +use ndarray::{Array, Axis, IxDyn}; + +use crate::{ops, Mask, MinOptMax, Options, OrtEngine, Y}; + +#[derive(Debug)] +pub struct MODNet { + engine: OrtEngine, + height: MinOptMax, + width: MinOptMax, + batch: MinOptMax, +} + +impl MODNet { + pub fn new(options: &Options) -> Result { + let engine = OrtEngine::new(options)?; + let (batch, height, width) = ( + engine.batch().to_owned(), + engine.height().to_owned(), + engine.width().to_owned(), + ); + engine.dry_run()?; + + Ok(Self { + engine, + height, + width, + batch, + }) + } + + pub fn run(&self, xs: &[DynamicImage]) -> Result> { + let xs_ = ops::resize(xs, self.height.opt as u32, self.width.opt as u32)?; + let xs_ = ops::normalize(xs_, 127.5, 255.0); + let ys = self.engine.run(&[xs_])?; + self.postprocess(ys, xs) + } + + pub fn postprocess(&self, xs: Vec>, xs0: &[DynamicImage]) -> Result> { + let mut ys: Vec = Vec::new(); + for (idx, luma) in xs[0].axis_iter(Axis(0)).enumerate() { + let luma = luma + .into_shape((self.height() as usize, self.width() as usize, 1))? + .into_owned(); + let v = luma + .into_raw_vec() + .iter() + .map(|x| (x * 255.0) as u8) + .collect::>(); + let luma: image::ImageBuffer, Vec> = + image::ImageBuffer::from_raw(self.width() as u32, self.height() as u32, v) + .expect("Faild to create image from ndarray"); + let luma = image::DynamicImage::from(luma); + let luma = luma.resize_exact( + xs0[idx].width(), + xs0[idx].height(), + image::imageops::FilterType::CatmullRom, + ); + ys.push(Y::default().with_masks(&[Mask::default().with_mask(luma)])); + } + Ok(ys) + } + + pub fn batch(&self) -> isize { + self.batch.opt + } + + pub fn width(&self) -> isize { + self.width.opt + } + + pub fn height(&self) -> isize { + self.height.opt + } +} diff --git a/src/models/yolo.rs b/src/models/yolo.rs index eb9102c..5cfedfe 100644 --- a/src/models/yolo.rs +++ b/src/models/yolo.rs @@ -4,7 +4,7 @@ use image::DynamicImage; use ndarray::{s, Array, Axis, IxDyn}; use regex::Regex; -use crate::{ops, Bbox, DynConf, Keypoint, Mask, Mbr, MinOptMax, Options, OrtEngine, Prob, Y}; +use crate::{ops, Bbox, DynConf, Keypoint, Mbr, MinOptMax, Options, OrtEngine, Polygon, Prob, Y}; const CXYWH_OFFSET: usize = 4; const KPT_STEP: usize = 3; @@ -313,7 +313,7 @@ impl YOLO { // masks if let YOLOTask::Segment = self.task { if let Some(bboxes) = y.bboxes() { - let mut y_masks: Vec = Vec::new(); + let mut y_polygons: Vec = Vec::new(); for bbox in bboxes.iter() { let coefs = if self.anchors_first { preds @@ -367,7 +367,7 @@ impl YOLO { } // get masks from image - let mut masks: Vec = Vec::new(); + let mut masks: Vec = Vec::new(); let contours: Vec> = imageproc::contours::find_contours_with_threshold( &mask_original, @@ -376,16 +376,16 @@ impl YOLO { contours.iter().for_each(|contour| { if contour.points.len() > 2 { masks.push( - Mask::default() + Polygon::default() .with_id(bbox.id()) .with_points_imageproc(&contour.points) .with_name(bbox.name().cloned()), ); } }); - y_masks.extend(masks); + y_polygons.extend(masks); } - y = y.with_masks(&y_masks); + y = y.with_polygons(&y_polygons); } } ys.push(y); diff --git a/src/models/yolop.rs b/src/models/yolop.rs index a08973c..81caa96 100644 --- a/src/models/yolop.rs +++ b/src/models/yolop.rs @@ -2,7 +2,7 @@ use anyhow::Result; use image::DynamicImage; use ndarray::{s, Array, Axis, IxDyn}; -use crate::{ops, Bbox, DynConf, Mask, MinOptMax, Options, OrtEngine, Y}; +use crate::{ops, Bbox, DynConf, MinOptMax, Options, OrtEngine, Polygon, Y}; #[derive(Debug)] pub struct YOLOPv2 { @@ -114,15 +114,15 @@ impl YOLOPv2 { image_height, ); let mask_da = mask_da.into_luma8(); - let mut y_masks: Vec = Vec::new(); + let mut y_polygons: Vec = Vec::new(); let contours: Vec> = imageproc::contours::find_contours_with_threshold(&mask_da, 1); contours.iter().for_each(|contour| { if contour.border_type == imageproc::contours::BorderType::Outer && contour.points.len() > 2 { - y_masks.push( - Mask::default() + y_polygons.push( + Polygon::default() .with_id(0) .with_points_imageproc(&contour.points) .with_name(Some("Drivable area".to_string())), @@ -151,26 +151,26 @@ impl YOLOPv2 { let mask_ll = mask_ll.into_luma8(); let contours: Vec> = imageproc::contours::find_contours_with_threshold(&mask_ll, 1); - let mut masks: Vec = Vec::new(); + let mut masks: Vec = Vec::new(); contours.iter().for_each(|contour| { if contour.border_type == imageproc::contours::BorderType::Outer && contour.points.len() > 2 { masks.push( - Mask::default() + Polygon::default() .with_id(1) .with_points_imageproc(&contour.points) .with_name(Some("Lane line".to_string())), ); } }); - y_masks.extend(masks); + y_polygons.extend(masks); // save ys.push( Y::default() .with_bboxes(&y_bboxes) - .with_masks(&y_masks) + .with_polygons(&y_polygons) .apply_bboxes_nms(self.iou), ); } diff --git a/src/utils/colormap256.rs b/src/utils/colormap256.rs new file mode 100644 index 0000000..96b8298 --- /dev/null +++ b/src/utils/colormap256.rs @@ -0,0 +1,2589 @@ +pub const TURBO: [[u8; 3]; 256] = [ + [48, 18, 59], + [50, 21, 67], + [51, 24, 74], + [52, 27, 81], + [53, 30, 88], + [54, 33, 95], + [55, 36, 102], + [56, 39, 109], + [57, 42, 115], + [58, 45, 121], + [59, 47, 128], + [60, 50, 134], + [61, 53, 139], + [62, 56, 145], + [63, 59, 151], + [63, 62, 156], + [64, 64, 162], + [65, 67, 167], + [65, 70, 172], + [66, 73, 177], + [66, 75, 181], + [67, 78, 186], + [68, 81, 191], + [68, 84, 195], + [68, 86, 199], + [69, 89, 203], + [69, 92, 207], + [69, 94, 211], + [70, 97, 214], + [70, 100, 218], + [70, 102, 221], + [70, 105, 224], + [70, 107, 227], + [71, 110, 230], + [71, 113, 233], + [71, 115, 235], + [71, 118, 238], + [71, 120, 240], + [71, 123, 242], + [70, 125, 244], + [70, 128, 246], + [70, 130, 248], + [70, 133, 250], + [70, 135, 251], + [69, 138, 252], + [69, 140, 253], + [68, 143, 254], + [67, 145, 254], + [66, 148, 255], + [65, 150, 255], + [64, 153, 255], + [62, 155, 254], + [61, 158, 254], + [59, 160, 253], + [58, 163, 252], + [56, 165, 251], + [55, 168, 250], + [53, 171, 248], + [51, 173, 247], + [49, 175, 245], + [47, 178, 244], + [46, 180, 242], + [44, 183, 240], + [42, 185, 238], + [40, 188, 235], + [39, 190, 233], + [37, 192, 231], + [35, 195, 228], + [34, 197, 226], + [32, 199, 223], + [31, 201, 221], + [30, 203, 218], + [28, 205, 216], + [27, 208, 213], + [26, 210, 210], + [26, 212, 208], + [25, 213, 205], + [24, 215, 202], + [24, 217, 200], + [24, 219, 197], + [24, 221, 194], + [24, 222, 192], + [24, 224, 189], + [25, 226, 187], + [25, 227, 185], + [26, 228, 182], + [28, 230, 180], + [29, 231, 178], + [31, 233, 175], + [32, 234, 172], + [34, 235, 170], + [37, 236, 167], + [39, 238, 164], + [42, 239, 161], + [44, 240, 158], + [47, 241, 155], + [50, 242, 152], + [53, 243, 148], + [56, 244, 145], + [60, 245, 142], + [63, 246, 138], + [67, 247, 135], + [70, 248, 132], + [74, 248, 128], + [78, 249, 125], + [82, 250, 122], + [85, 250, 118], + [89, 251, 115], + [93, 252, 111], + [97, 252, 108], + [101, 253, 105], + [105, 253, 102], + [109, 254, 98], + [113, 254, 95], + [117, 254, 92], + [121, 254, 89], + [125, 255, 86], + [128, 255, 83], + [132, 255, 81], + [136, 255, 78], + [139, 255, 75], + [143, 255, 73], + [146, 255, 71], + [150, 254, 68], + [153, 254, 66], + [156, 254, 64], + [159, 253, 63], + [161, 253, 61], + [164, 252, 60], + [167, 252, 58], + [169, 251, 57], + [172, 251, 56], + [175, 250, 55], + [177, 249, 54], + [180, 248, 54], + [183, 247, 53], + [185, 246, 53], + [188, 245, 52], + [190, 244, 52], + [193, 243, 52], + [195, 241, 52], + [198, 240, 52], + [200, 239, 52], + [203, 237, 52], + [205, 236, 52], + [208, 234, 52], + [210, 233, 53], + [212, 231, 53], + [215, 229, 53], + [217, 228, 54], + [219, 226, 54], + [221, 224, 55], + [223, 223, 55], + [225, 221, 55], + [227, 219, 56], + [229, 217, 56], + [231, 215, 57], + [233, 213, 57], + [235, 211, 57], + [236, 209, 58], + [238, 207, 58], + [239, 205, 58], + [241, 203, 58], + [242, 201, 58], + [244, 199, 58], + [245, 197, 58], + [246, 195, 58], + [247, 193, 58], + [248, 190, 57], + [249, 188, 57], + [250, 186, 57], + [251, 184, 56], + [251, 182, 55], + [252, 179, 54], + [252, 177, 54], + [253, 174, 53], + [253, 172, 52], + [254, 169, 51], + [254, 167, 50], + [254, 164, 49], + [254, 161, 48], + [254, 158, 47], + [254, 155, 45], + [254, 153, 44], + [254, 150, 43], + [254, 147, 42], + [254, 144, 41], + [253, 141, 39], + [253, 138, 38], + [252, 135, 37], + [252, 132, 35], + [251, 129, 34], + [251, 126, 33], + [250, 123, 31], + [249, 120, 30], + [249, 117, 29], + [248, 114, 28], + [247, 111, 26], + [246, 108, 25], + [245, 105, 24], + [244, 102, 23], + [243, 99, 21], + [242, 96, 20], + [241, 93, 19], + [240, 91, 18], + [239, 88, 17], + [237, 85, 16], + [236, 83, 15], + [235, 80, 14], + [234, 78, 13], + [232, 75, 12], + [231, 73, 12], + [229, 71, 11], + [228, 69, 10], + [226, 67, 10], + [225, 65, 9], + [223, 63, 8], + [221, 61, 8], + [220, 59, 7], + [218, 57, 7], + [216, 55, 6], + [214, 53, 6], + [212, 51, 5], + [210, 49, 5], + [208, 47, 5], + [206, 45, 4], + [204, 43, 4], + [202, 42, 4], + [200, 40, 3], + [197, 38, 3], + [195, 37, 3], + [193, 35, 2], + [190, 33, 2], + [188, 32, 2], + [185, 30, 2], + [183, 29, 2], + [180, 27, 1], + [178, 26, 1], + [175, 24, 1], + [172, 23, 1], + [169, 22, 1], + [167, 20, 1], + [164, 19, 1], + [161, 18, 1], + [158, 16, 1], + [155, 15, 1], + [152, 14, 1], + [149, 13, 1], + [146, 11, 1], + [142, 10, 1], + [139, 9, 2], + [136, 8, 2], + [133, 7, 2], + [129, 6, 2], + [126, 5, 2], + [122, 4, 3], +]; + +pub const INFERNO: [[u8; 3]; 256] = [ + [0, 0, 4], + [1, 0, 5], + [1, 1, 6], + [1, 1, 8], + [2, 1, 10], + [2, 2, 12], + [2, 2, 14], + [3, 2, 16], + [4, 3, 18], + [4, 3, 20], + [5, 4, 23], + [6, 4, 25], + [7, 5, 27], + [8, 5, 29], + [9, 6, 31], + [10, 7, 34], + [11, 7, 36], + [12, 8, 38], + [13, 8, 41], + [14, 9, 43], + [16, 9, 45], + [17, 10, 48], + [18, 10, 50], + [20, 11, 52], + [21, 11, 55], + [22, 11, 57], + [24, 12, 60], + [25, 12, 62], + [27, 12, 65], + [28, 12, 67], + [30, 12, 69], + [31, 12, 72], + [33, 12, 74], + [35, 12, 76], + [36, 12, 79], + [38, 12, 81], + [40, 11, 83], + [41, 11, 85], + [43, 11, 87], + [45, 11, 89], + [47, 10, 91], + [49, 10, 92], + [50, 10, 94], + [52, 10, 95], + [54, 9, 97], + [56, 9, 98], + [57, 9, 99], + [59, 9, 100], + [61, 9, 101], + [62, 9, 102], + [64, 10, 103], + [66, 10, 104], + [68, 10, 104], + [69, 10, 105], + [71, 11, 106], + [73, 11, 106], + [74, 12, 107], + [76, 12, 107], + [77, 13, 108], + [79, 13, 108], + [81, 14, 108], + [82, 14, 109], + [84, 15, 109], + [85, 15, 109], + [87, 16, 110], + [89, 16, 110], + [90, 17, 110], + [92, 18, 110], + [93, 18, 110], + [95, 19, 110], + [97, 19, 110], + [98, 20, 110], + [100, 21, 110], + [101, 21, 110], + [103, 22, 110], + [105, 22, 110], + [106, 23, 110], + [108, 24, 110], + [109, 24, 110], + [111, 25, 110], + [113, 25, 110], + [114, 26, 110], + [116, 26, 110], + [117, 27, 110], + [119, 28, 109], + [120, 28, 109], + [122, 29, 109], + [124, 29, 109], + [125, 30, 109], + [127, 30, 108], + [128, 31, 108], + [130, 32, 108], + [132, 32, 107], + [133, 33, 107], + [135, 33, 107], + [136, 34, 106], + [138, 34, 106], + [140, 35, 105], + [141, 35, 105], + [143, 36, 105], + [144, 37, 104], + [146, 37, 104], + [147, 38, 103], + [149, 38, 103], + [151, 39, 102], + [152, 39, 102], + [154, 40, 101], + [155, 41, 100], + [157, 41, 100], + [159, 42, 99], + [160, 42, 99], + [162, 43, 98], + [163, 44, 97], + [165, 44, 96], + [166, 45, 96], + [168, 46, 95], + [169, 46, 94], + [171, 47, 94], + [173, 48, 93], + [174, 48, 92], + [176, 49, 91], + [177, 50, 90], + [179, 50, 90], + [180, 51, 89], + [182, 52, 88], + [183, 53, 87], + [185, 53, 86], + [186, 54, 85], + [188, 55, 84], + [189, 56, 83], + [191, 57, 82], + [192, 58, 81], + [193, 58, 80], + [195, 59, 79], + [196, 60, 78], + [198, 61, 77], + [199, 62, 76], + [200, 63, 75], + [202, 64, 74], + [203, 65, 73], + [204, 66, 72], + [206, 67, 71], + [207, 68, 70], + [208, 69, 69], + [210, 70, 68], + [211, 71, 67], + [212, 72, 66], + [213, 74, 65], + [215, 75, 63], + [216, 76, 62], + [217, 77, 61], + [218, 78, 60], + [219, 80, 59], + [221, 81, 58], + [222, 82, 56], + [223, 83, 55], + [224, 85, 54], + [225, 86, 53], + [226, 87, 52], + [227, 89, 51], + [228, 90, 49], + [229, 92, 48], + [230, 93, 47], + [231, 94, 46], + [232, 96, 45], + [233, 97, 43], + [234, 99, 42], + [235, 100, 41], + [235, 102, 40], + [236, 103, 38], + [237, 105, 37], + [238, 106, 36], + [239, 108, 35], + [239, 110, 33], + [240, 111, 32], + [241, 113, 31], + [241, 115, 29], + [242, 116, 28], + [243, 118, 27], + [243, 120, 25], + [244, 121, 24], + [245, 123, 23], + [245, 125, 21], + [246, 126, 20], + [246, 128, 19], + [247, 130, 18], + [247, 132, 16], + [248, 133, 15], + [248, 135, 14], + [248, 137, 12], + [249, 139, 11], + [249, 140, 10], + [249, 142, 9], + [250, 144, 8], + [250, 146, 7], + [250, 148, 7], + [251, 150, 6], + [251, 151, 6], + [251, 153, 6], + [251, 155, 6], + [251, 157, 7], + [252, 159, 7], + [252, 161, 8], + [252, 163, 9], + [252, 165, 10], + [252, 166, 12], + [252, 168, 13], + [252, 170, 15], + [252, 172, 17], + [252, 174, 18], + [252, 176, 20], + [252, 178, 22], + [252, 180, 24], + [251, 182, 26], + [251, 184, 29], + [251, 186, 31], + [251, 188, 33], + [251, 190, 35], + [250, 192, 38], + [250, 194, 40], + [250, 196, 42], + [250, 198, 45], + [249, 199, 47], + [249, 201, 50], + [249, 203, 53], + [248, 205, 55], + [248, 207, 58], + [247, 209, 61], + [247, 211, 64], + [246, 213, 67], + [246, 215, 70], + [245, 217, 73], + [245, 219, 76], + [244, 221, 79], + [244, 223, 83], + [244, 225, 86], + [243, 227, 90], + [243, 229, 93], + [242, 230, 97], + [242, 232, 101], + [242, 234, 105], + [241, 236, 109], + [241, 237, 113], + [241, 239, 117], + [241, 241, 121], + [242, 242, 125], + [242, 244, 130], + [243, 245, 134], + [243, 246, 138], + [244, 248, 142], + [245, 249, 146], + [246, 250, 150], + [248, 251, 154], + [249, 252, 157], + [250, 253, 161], + [252, 255, 164], +]; + +pub const PLASMA: [[u8; 3]; 256] = [ + [13, 8, 135], + [16, 7, 136], + [19, 7, 137], + [22, 7, 138], + [25, 6, 140], + [27, 6, 141], + [29, 6, 142], + [32, 6, 143], + [34, 6, 144], + [36, 6, 145], + [38, 5, 145], + [40, 5, 146], + [42, 5, 147], + [44, 5, 148], + [46, 5, 149], + [47, 5, 150], + [49, 5, 151], + [51, 5, 151], + [53, 4, 152], + [55, 4, 153], + [56, 4, 154], + [58, 4, 154], + [60, 4, 155], + [62, 4, 156], + [63, 4, 156], + [65, 4, 157], + [67, 3, 158], + [68, 3, 158], + [70, 3, 159], + [72, 3, 159], + [73, 3, 160], + [75, 3, 161], + [76, 2, 161], + [78, 2, 162], + [80, 2, 162], + [81, 2, 163], + [83, 2, 163], + [85, 2, 164], + [86, 1, 164], + [88, 1, 164], + [89, 1, 165], + [91, 1, 165], + [92, 1, 166], + [94, 1, 166], + [96, 1, 166], + [97, 0, 167], + [99, 0, 167], + [100, 0, 167], + [102, 0, 167], + [103, 0, 168], + [105, 0, 168], + [106, 0, 168], + [108, 0, 168], + [110, 0, 168], + [111, 0, 168], + [113, 0, 168], + [114, 1, 168], + [116, 1, 168], + [117, 1, 168], + [119, 1, 168], + [120, 1, 168], + [122, 2, 168], + [123, 2, 168], + [125, 3, 168], + [126, 3, 168], + [128, 4, 168], + [129, 4, 167], + [131, 5, 167], + [132, 5, 167], + [134, 6, 166], + [135, 7, 166], + [136, 8, 166], + [138, 9, 165], + [139, 10, 165], + [141, 11, 165], + [142, 12, 164], + [143, 13, 164], + [145, 14, 163], + [146, 15, 163], + [148, 16, 162], + [149, 17, 161], + [150, 19, 161], + [152, 20, 160], + [153, 21, 159], + [154, 22, 159], + [156, 23, 158], + [157, 24, 157], + [158, 25, 157], + [160, 26, 156], + [161, 27, 155], + [162, 29, 154], + [163, 30, 154], + [165, 31, 153], + [166, 32, 152], + [167, 33, 151], + [168, 34, 150], + [170, 35, 149], + [171, 36, 148], + [172, 38, 148], + [173, 39, 147], + [174, 40, 146], + [176, 41, 145], + [177, 42, 144], + [178, 43, 143], + [179, 44, 142], + [180, 46, 141], + [181, 47, 140], + [182, 48, 139], + [183, 49, 138], + [184, 50, 137], + [186, 51, 136], + [187, 52, 136], + [188, 53, 135], + [189, 55, 134], + [190, 56, 133], + [191, 57, 132], + [192, 58, 131], + [193, 59, 130], + [194, 60, 129], + [195, 61, 128], + [196, 62, 127], + [197, 64, 126], + [198, 65, 125], + [199, 66, 124], + [200, 67, 123], + [201, 68, 122], + [202, 69, 122], + [203, 70, 121], + [204, 71, 120], + [204, 73, 119], + [205, 74, 118], + [206, 75, 117], + [207, 76, 116], + [208, 77, 115], + [209, 78, 114], + [210, 79, 113], + [211, 81, 113], + [212, 82, 112], + [213, 83, 111], + [213, 84, 110], + [214, 85, 109], + [215, 86, 108], + [216, 87, 107], + [217, 88, 106], + [218, 90, 106], + [218, 91, 105], + [219, 92, 104], + [220, 93, 103], + [221, 94, 102], + [222, 95, 101], + [222, 97, 100], + [223, 98, 99], + [224, 99, 99], + [225, 100, 98], + [226, 101, 97], + [226, 102, 96], + [227, 104, 95], + [228, 105, 94], + [229, 106, 93], + [229, 107, 93], + [230, 108, 92], + [231, 110, 91], + [231, 111, 90], + [232, 112, 89], + [233, 113, 88], + [233, 114, 87], + [234, 116, 87], + [235, 117, 86], + [235, 118, 85], + [236, 119, 84], + [237, 121, 83], + [237, 122, 82], + [238, 123, 81], + [239, 124, 81], + [239, 126, 80], + [240, 127, 79], + [240, 128, 78], + [241, 129, 77], + [241, 131, 76], + [242, 132, 75], + [243, 133, 75], + [243, 135, 74], + [244, 136, 73], + [244, 137, 72], + [245, 139, 71], + [245, 140, 70], + [246, 141, 69], + [246, 143, 68], + [247, 144, 68], + [247, 145, 67], + [247, 147, 66], + [248, 148, 65], + [248, 149, 64], + [249, 151, 63], + [249, 152, 62], + [249, 154, 62], + [250, 155, 61], + [250, 156, 60], + [250, 158, 59], + [251, 159, 58], + [251, 161, 57], + [251, 162, 56], + [252, 163, 56], + [252, 165, 55], + [252, 166, 54], + [252, 168, 53], + [252, 169, 52], + [253, 171, 51], + [253, 172, 51], + [253, 174, 50], + [253, 175, 49], + [253, 177, 48], + [253, 178, 47], + [253, 180, 47], + [253, 181, 46], + [254, 183, 45], + [254, 184, 44], + [254, 186, 44], + [254, 187, 43], + [254, 189, 42], + [254, 190, 42], + [254, 192, 41], + [253, 194, 41], + [253, 195, 40], + [253, 197, 39], + [253, 198, 39], + [253, 200, 39], + [253, 202, 38], + [253, 203, 38], + [252, 205, 37], + [252, 206, 37], + [252, 208, 37], + [252, 210, 37], + [251, 211, 36], + [251, 213, 36], + [251, 215, 36], + [250, 216, 36], + [250, 218, 36], + [249, 220, 36], + [249, 221, 37], + [248, 223, 37], + [248, 225, 37], + [247, 226, 37], + [247, 228, 37], + [246, 230, 38], + [246, 232, 38], + [245, 233, 38], + [245, 235, 39], + [244, 237, 39], + [243, 238, 39], + [243, 240, 39], + [242, 242, 39], + [241, 244, 38], + [241, 245, 37], + [240, 247, 36], + [240, 249, 33], +]; + +pub const VIRIDIS: [[u8; 3]; 256] = [ + [68, 1, 84], + [68, 2, 86], + [69, 4, 87], + [69, 5, 89], + [70, 7, 90], + [70, 8, 92], + [70, 10, 93], + [70, 11, 94], + [71, 13, 96], + [71, 14, 97], + [71, 16, 99], + [71, 17, 100], + [71, 19, 101], + [72, 20, 103], + [72, 22, 104], + [72, 23, 105], + [72, 24, 106], + [72, 26, 108], + [72, 27, 109], + [72, 28, 110], + [72, 29, 111], + [72, 31, 112], + [72, 32, 113], + [72, 33, 115], + [72, 35, 116], + [72, 36, 117], + [72, 37, 118], + [72, 38, 119], + [72, 40, 120], + [72, 41, 121], + [71, 42, 122], + [71, 44, 122], + [71, 45, 123], + [71, 46, 124], + [71, 47, 125], + [70, 48, 126], + [70, 50, 126], + [70, 51, 127], + [70, 52, 128], + [69, 53, 129], + [69, 55, 129], + [69, 56, 130], + [68, 57, 131], + [68, 58, 131], + [68, 59, 132], + [67, 61, 132], + [67, 62, 133], + [66, 63, 133], + [66, 64, 134], + [66, 65, 134], + [65, 66, 135], + [65, 68, 135], + [64, 69, 136], + [64, 70, 136], + [63, 71, 136], + [63, 72, 137], + [62, 73, 137], + [62, 74, 137], + [62, 76, 138], + [61, 77, 138], + [61, 78, 138], + [60, 79, 138], + [60, 80, 139], + [59, 81, 139], + [59, 82, 139], + [58, 83, 139], + [58, 84, 140], + [57, 85, 140], + [57, 86, 140], + [56, 88, 140], + [56, 89, 140], + [55, 90, 140], + [55, 91, 141], + [54, 92, 141], + [54, 93, 141], + [53, 94, 141], + [53, 95, 141], + [52, 96, 141], + [52, 97, 141], + [51, 98, 141], + [51, 99, 141], + [50, 100, 142], + [50, 101, 142], + [49, 102, 142], + [49, 103, 142], + [49, 104, 142], + [48, 105, 142], + [48, 106, 142], + [47, 107, 142], + [47, 108, 142], + [46, 109, 142], + [46, 110, 142], + [46, 111, 142], + [45, 112, 142], + [45, 113, 142], + [44, 113, 142], + [44, 114, 142], + [44, 115, 142], + [43, 116, 142], + [43, 117, 142], + [42, 118, 142], + [42, 119, 142], + [42, 120, 142], + [41, 121, 142], + [41, 122, 142], + [41, 123, 142], + [40, 124, 142], + [40, 125, 142], + [39, 126, 142], + [39, 127, 142], + [39, 128, 142], + [38, 129, 142], + [38, 130, 142], + [38, 130, 142], + [37, 131, 142], + [37, 132, 142], + [37, 133, 142], + [36, 134, 142], + [36, 135, 142], + [35, 136, 142], + [35, 137, 142], + [35, 138, 141], + [34, 139, 141], + [34, 140, 141], + [34, 141, 141], + [33, 142, 141], + [33, 143, 141], + [33, 144, 141], + [33, 145, 140], + [32, 146, 140], + [32, 146, 140], + [32, 147, 140], + [31, 148, 140], + [31, 149, 139], + [31, 150, 139], + [31, 151, 139], + [31, 152, 139], + [31, 153, 138], + [31, 154, 138], + [30, 155, 138], + [30, 156, 137], + [30, 157, 137], + [31, 158, 137], + [31, 159, 136], + [31, 160, 136], + [31, 161, 136], + [31, 161, 135], + [31, 162, 135], + [32, 163, 134], + [32, 164, 134], + [33, 165, 133], + [33, 166, 133], + [34, 167, 133], + [34, 168, 132], + [35, 169, 131], + [36, 170, 131], + [37, 171, 130], + [37, 172, 130], + [38, 173, 129], + [39, 173, 129], + [40, 174, 128], + [41, 175, 127], + [42, 176, 127], + [44, 177, 126], + [45, 178, 125], + [46, 179, 124], + [47, 180, 124], + [49, 181, 123], + [50, 182, 122], + [52, 182, 121], + [53, 183, 121], + [55, 184, 120], + [56, 185, 119], + [58, 186, 118], + [59, 187, 117], + [61, 188, 116], + [63, 188, 115], + [64, 189, 114], + [66, 190, 113], + [68, 191, 112], + [70, 192, 111], + [72, 193, 110], + [74, 193, 109], + [76, 194, 108], + [78, 195, 107], + [80, 196, 106], + [82, 197, 105], + [84, 197, 104], + [86, 198, 103], + [88, 199, 101], + [90, 200, 100], + [92, 200, 99], + [94, 201, 98], + [96, 202, 96], + [99, 203, 95], + [101, 203, 94], + [103, 204, 92], + [105, 205, 91], + [108, 205, 90], + [110, 206, 88], + [112, 207, 87], + [115, 208, 86], + [117, 208, 84], + [119, 209, 83], + [122, 209, 81], + [124, 210, 80], + [127, 211, 78], + [129, 211, 77], + [132, 212, 75], + [134, 213, 73], + [137, 213, 72], + [139, 214, 70], + [142, 214, 69], + [144, 215, 67], + [147, 215, 65], + [149, 216, 64], + [152, 216, 62], + [155, 217, 60], + [157, 217, 59], + [160, 218, 57], + [162, 218, 55], + [165, 219, 54], + [168, 219, 52], + [170, 220, 50], + [173, 220, 48], + [176, 221, 47], + [178, 221, 45], + [181, 222, 43], + [184, 222, 41], + [186, 222, 40], + [189, 223, 38], + [192, 223, 37], + [194, 223, 35], + [197, 224, 33], + [200, 224, 32], + [202, 225, 31], + [205, 225, 29], + [208, 225, 28], + [210, 226, 27], + [213, 226, 26], + [216, 226, 25], + [218, 227, 25], + [221, 227, 24], + [223, 227, 24], + [226, 228, 24], + [229, 228, 25], + [231, 228, 25], + [234, 229, 26], + [236, 229, 27], + [239, 229, 28], + [241, 229, 29], + [244, 230, 30], + [246, 230, 32], + [248, 230, 33], + [251, 231, 35], + [253, 231, 37], +]; + +pub const MAGMA: [[u8; 3]; 256] = [ + [0, 0, 4], + [1, 0, 5], + [1, 1, 6], + [1, 1, 8], + [2, 1, 9], + [2, 2, 11], + [2, 2, 13], + [3, 3, 15], + [3, 3, 18], + [4, 4, 20], + [5, 4, 22], + [6, 5, 24], + [6, 5, 26], + [7, 6, 28], + [8, 7, 30], + [9, 7, 32], + [10, 8, 34], + [11, 9, 36], + [12, 9, 38], + [13, 10, 41], + [14, 11, 43], + [16, 11, 45], + [17, 12, 47], + [18, 13, 49], + [19, 13, 52], + [20, 14, 54], + [21, 14, 56], + [22, 15, 59], + [24, 15, 61], + [25, 16, 63], + [26, 16, 66], + [28, 16, 68], + [29, 17, 71], + [30, 17, 73], + [32, 17, 75], + [33, 17, 78], + [34, 17, 80], + [36, 18, 83], + [37, 18, 85], + [39, 18, 88], + [41, 17, 90], + [42, 17, 92], + [44, 17, 95], + [45, 17, 97], + [47, 17, 99], + [49, 17, 101], + [51, 16, 103], + [52, 16, 105], + [54, 16, 107], + [56, 16, 108], + [57, 15, 110], + [59, 15, 112], + [61, 15, 113], + [63, 15, 114], + [64, 15, 116], + [66, 15, 117], + [68, 15, 118], + [69, 16, 119], + [71, 16, 120], + [73, 16, 120], + [74, 16, 121], + [76, 17, 122], + [78, 17, 123], + [79, 18, 123], + [81, 18, 124], + [82, 19, 124], + [84, 19, 125], + [86, 20, 125], + [87, 21, 126], + [89, 21, 126], + [90, 22, 126], + [92, 22, 127], + [93, 23, 127], + [95, 24, 127], + [96, 24, 128], + [98, 25, 128], + [100, 26, 128], + [101, 26, 128], + [103, 27, 128], + [104, 28, 129], + [106, 28, 129], + [107, 29, 129], + [109, 29, 129], + [110, 30, 129], + [112, 31, 129], + [114, 31, 129], + [115, 32, 129], + [117, 33, 129], + [118, 33, 129], + [120, 34, 129], + [121, 34, 130], + [123, 35, 130], + [124, 35, 130], + [126, 36, 130], + [128, 37, 130], + [129, 37, 129], + [131, 38, 129], + [132, 38, 129], + [134, 39, 129], + [136, 39, 129], + [137, 40, 129], + [139, 41, 129], + [140, 41, 129], + [142, 42, 129], + [144, 42, 129], + [145, 43, 129], + [147, 43, 128], + [148, 44, 128], + [150, 44, 128], + [152, 45, 128], + [153, 45, 128], + [155, 46, 127], + [156, 46, 127], + [158, 47, 127], + [160, 47, 127], + [161, 48, 126], + [163, 48, 126], + [165, 49, 126], + [166, 49, 125], + [168, 50, 125], + [170, 51, 125], + [171, 51, 124], + [173, 52, 124], + [174, 52, 123], + [176, 53, 123], + [178, 53, 123], + [179, 54, 122], + [181, 54, 122], + [183, 55, 121], + [184, 55, 121], + [186, 56, 120], + [188, 57, 120], + [189, 57, 119], + [191, 58, 119], + [192, 58, 118], + [194, 59, 117], + [196, 60, 117], + [197, 60, 116], + [199, 61, 115], + [200, 62, 115], + [202, 62, 114], + [204, 63, 113], + [205, 64, 113], + [207, 64, 112], + [208, 65, 111], + [210, 66, 111], + [211, 67, 110], + [213, 68, 109], + [214, 69, 108], + [216, 69, 108], + [217, 70, 107], + [219, 71, 106], + [220, 72, 105], + [222, 73, 104], + [223, 74, 104], + [224, 76, 103], + [226, 77, 102], + [227, 78, 101], + [228, 79, 100], + [229, 80, 100], + [231, 82, 99], + [232, 83, 98], + [233, 84, 98], + [234, 86, 97], + [235, 87, 96], + [236, 88, 96], + [237, 90, 95], + [238, 91, 94], + [239, 93, 94], + [240, 95, 94], + [241, 96, 93], + [242, 98, 93], + [242, 100, 92], + [243, 101, 92], + [244, 103, 92], + [244, 105, 92], + [245, 107, 92], + [246, 108, 92], + [246, 110, 92], + [247, 112, 92], + [247, 114, 92], + [248, 116, 92], + [248, 118, 92], + [249, 120, 93], + [249, 121, 93], + [249, 123, 93], + [250, 125, 94], + [250, 127, 94], + [250, 129, 95], + [251, 131, 95], + [251, 133, 96], + [251, 135, 97], + [252, 137, 97], + [252, 138, 98], + [252, 140, 99], + [252, 142, 100], + [252, 144, 101], + [253, 146, 102], + [253, 148, 103], + [253, 150, 104], + [253, 152, 105], + [253, 154, 106], + [253, 155, 107], + [254, 157, 108], + [254, 159, 109], + [254, 161, 110], + [254, 163, 111], + [254, 165, 113], + [254, 167, 114], + [254, 169, 115], + [254, 170, 116], + [254, 172, 118], + [254, 174, 119], + [254, 176, 120], + [254, 178, 122], + [254, 180, 123], + [254, 182, 124], + [254, 183, 126], + [254, 185, 127], + [254, 187, 129], + [254, 189, 130], + [254, 191, 132], + [254, 193, 133], + [254, 194, 135], + [254, 196, 136], + [254, 198, 138], + [254, 200, 140], + [254, 202, 141], + [254, 204, 143], + [254, 205, 144], + [254, 207, 146], + [254, 209, 148], + [254, 211, 149], + [254, 213, 151], + [254, 215, 153], + [254, 216, 154], + [253, 218, 156], + [253, 220, 158], + [253, 222, 160], + [253, 224, 161], + [253, 226, 163], + [253, 227, 165], + [253, 229, 167], + [253, 231, 169], + [253, 233, 170], + [253, 235, 172], + [252, 236, 174], + [252, 238, 176], + [252, 240, 178], + [252, 242, 180], + [252, 244, 182], + [252, 246, 184], + [252, 247, 185], + [252, 249, 187], + [252, 251, 189], + [252, 253, 191], +]; + +pub const BENTCOOLWARM: [[u8; 3]; 256] = [ + [59, 76, 192], + [60, 78, 193], + [61, 79, 194], + [62, 80, 194], + [63, 82, 195], + [64, 83, 196], + [65, 85, 196], + [66, 86, 197], + [67, 87, 197], + [68, 89, 198], + [69, 90, 199], + [70, 91, 199], + [71, 93, 200], + [72, 94, 200], + [73, 95, 201], + [74, 97, 202], + [75, 98, 202], + [76, 100, 203], + [77, 101, 203], + [78, 102, 204], + [79, 104, 204], + [81, 105, 205], + [82, 106, 206], + [83, 108, 206], + [84, 109, 207], + [85, 110, 207], + [86, 112, 208], + [87, 113, 208], + [89, 114, 209], + [90, 116, 209], + [91, 117, 210], + [92, 118, 210], + [93, 120, 211], + [95, 121, 211], + [96, 122, 212], + [97, 124, 212], + [98, 125, 213], + [100, 126, 213], + [101, 128, 214], + [102, 129, 214], + [103, 130, 214], + [105, 131, 215], + [106, 133, 215], + [107, 134, 216], + [109, 135, 216], + [110, 137, 217], + [111, 138, 217], + [113, 139, 218], + [114, 141, 218], + [115, 142, 218], + [117, 143, 219], + [118, 145, 219], + [119, 146, 220], + [121, 147, 220], + [122, 149, 220], + [123, 150, 221], + [125, 151, 221], + [126, 152, 222], + [128, 154, 222], + [129, 155, 222], + [130, 156, 223], + [132, 158, 223], + [133, 159, 223], + [135, 160, 224], + [136, 162, 224], + [138, 163, 225], + [139, 164, 225], + [140, 165, 225], + [142, 167, 226], + [143, 168, 226], + [145, 169, 226], + [146, 171, 227], + [148, 172, 227], + [149, 173, 227], + [151, 175, 228], + [152, 176, 228], + [154, 177, 228], + [155, 178, 229], + [157, 180, 229], + [159, 181, 229], + [160, 182, 230], + [162, 184, 230], + [163, 185, 230], + [165, 186, 230], + [166, 187, 231], + [168, 189, 231], + [170, 190, 231], + [171, 191, 232], + [173, 193, 232], + [174, 194, 232], + [176, 195, 232], + [178, 196, 233], + [179, 198, 233], + [181, 199, 233], + [183, 200, 234], + [184, 202, 234], + [186, 203, 234], + [188, 204, 234], + [189, 205, 235], + [191, 207, 235], + [193, 208, 235], + [194, 209, 236], + [196, 210, 236], + [198, 212, 236], + [200, 213, 236], + [201, 214, 237], + [203, 215, 237], + [205, 217, 237], + [207, 218, 237], + [208, 219, 238], + [210, 220, 238], + [212, 222, 238], + [214, 223, 238], + [215, 224, 239], + [217, 225, 239], + [219, 227, 239], + [221, 228, 239], + [223, 229, 240], + [225, 230, 240], + [226, 232, 240], + [228, 233, 240], + [230, 234, 241], + [232, 235, 241], + [234, 237, 241], + [236, 238, 241], + [238, 239, 242], + [239, 240, 242], + [241, 242, 242], + [242, 241, 241], + [242, 240, 239], + [241, 238, 237], + [241, 237, 235], + [241, 235, 232], + [241, 234, 230], + [240, 232, 228], + [240, 231, 226], + [240, 229, 224], + [239, 228, 222], + [239, 226, 219], + [239, 225, 217], + [238, 223, 215], + [238, 222, 213], + [238, 220, 211], + [237, 219, 209], + [237, 217, 207], + [237, 216, 205], + [236, 214, 203], + [236, 213, 201], + [236, 211, 199], + [235, 210, 196], + [235, 208, 194], + [235, 207, 192], + [234, 205, 190], + [234, 204, 188], + [233, 202, 186], + [233, 201, 184], + [233, 199, 182], + [232, 197, 180], + [232, 196, 179], + [232, 194, 177], + [231, 193, 175], + [231, 191, 173], + [230, 190, 171], + [230, 188, 169], + [230, 187, 167], + [229, 185, 165], + [229, 184, 163], + [228, 182, 161], + [228, 181, 159], + [228, 179, 158], + [227, 177, 156], + [227, 176, 154], + [226, 174, 152], + [226, 173, 150], + [226, 171, 148], + [225, 170, 147], + [225, 168, 145], + [224, 167, 143], + [224, 165, 141], + [223, 163, 140], + [223, 162, 138], + [223, 160, 136], + [222, 159, 134], + [222, 157, 133], + [221, 156, 131], + [221, 154, 129], + [220, 152, 128], + [220, 151, 126], + [219, 149, 124], + [219, 148, 123], + [218, 146, 121], + [218, 144, 119], + [217, 143, 118], + [217, 141, 116], + [217, 140, 114], + [216, 138, 113], + [216, 136, 111], + [215, 135, 110], + [215, 133, 108], + [214, 132, 107], + [214, 130, 105], + [213, 128, 104], + [212, 127, 102], + [212, 125, 101], + [211, 123, 99], + [211, 122, 98], + [210, 120, 96], + [210, 119, 95], + [209, 117, 93], + [209, 115, 92], + [208, 114, 90], + [208, 112, 89], + [207, 110, 88], + [207, 108, 86], + [206, 107, 85], + [205, 105, 83], + [205, 103, 82], + [204, 102, 81], + [204, 100, 79], + [203, 98, 78], + [203, 96, 77], + [202, 95, 75], + [201, 93, 74], + [201, 91, 73], + [200, 89, 72], + [200, 87, 70], + [199, 86, 69], + [198, 84, 68], + [198, 82, 67], + [197, 80, 65], + [197, 78, 64], + [196, 76, 63], + [195, 74, 62], + [195, 72, 61], + [194, 70, 60], + [193, 68, 58], + [193, 66, 57], + [192, 64, 56], + [192, 62, 55], + [191, 60, 54], + [190, 58, 53], + [190, 55, 52], + [189, 53, 51], + [188, 50, 50], + [188, 48, 49], + [187, 45, 48], + [186, 43, 47], + [186, 40, 46], + [185, 37, 45], + [184, 34, 44], + [184, 30, 43], + [183, 27, 42], + [182, 22, 41], + [181, 18, 40], + [181, 12, 39], + [180, 4, 38], +]; + +pub const BLACKBODY: [[u8; 3]; 256] = [ + [0, 0, 0], + [3, 1, 1], + [7, 2, 1], + [10, 3, 2], + [13, 4, 2], + [16, 5, 3], + [18, 6, 3], + [20, 7, 4], + [22, 8, 4], + [24, 9, 5], + [26, 10, 5], + [27, 11, 6], + [29, 11, 6], + [30, 12, 7], + [32, 13, 8], + [33, 14, 8], + [34, 15, 9], + [36, 15, 9], + [37, 16, 10], + [38, 16, 10], + [40, 17, 11], + [41, 17, 11], + [43, 18, 12], + [44, 18, 12], + [46, 18, 13], + [47, 19, 13], + [49, 19, 14], + [50, 19, 14], + [52, 20, 15], + [54, 20, 15], + [55, 20, 15], + [57, 21, 16], + [58, 21, 16], + [60, 21, 16], + [62, 22, 17], + [63, 22, 17], + [65, 22, 17], + [66, 23, 18], + [68, 23, 18], + [70, 23, 18], + [71, 24, 19], + [73, 24, 19], + [75, 24, 19], + [76, 25, 20], + [78, 25, 20], + [80, 25, 20], + [81, 25, 20], + [83, 26, 21], + [85, 26, 21], + [86, 26, 21], + [88, 26, 21], + [90, 27, 22], + [91, 27, 22], + [93, 27, 22], + [95, 27, 22], + [97, 28, 23], + [98, 28, 23], + [100, 28, 23], + [102, 28, 23], + [104, 29, 24], + [105, 29, 24], + [107, 29, 24], + [109, 29, 24], + [111, 29, 25], + [112, 30, 25], + [114, 30, 25], + [116, 30, 25], + [118, 30, 26], + [119, 30, 26], + [121, 31, 26], + [123, 31, 26], + [125, 31, 27], + [127, 31, 27], + [128, 31, 27], + [130, 31, 27], + [132, 32, 28], + [134, 32, 28], + [136, 32, 28], + [137, 32, 28], + [139, 32, 29], + [141, 32, 29], + [143, 32, 29], + [145, 33, 29], + [147, 33, 30], + [148, 33, 30], + [150, 33, 30], + [152, 33, 31], + [154, 33, 31], + [156, 33, 31], + [158, 33, 31], + [160, 33, 32], + [161, 34, 32], + [163, 34, 32], + [165, 34, 32], + [167, 34, 33], + [169, 34, 33], + [171, 34, 33], + [173, 34, 33], + [175, 34, 34], + [177, 34, 34], + [178, 34, 34], + [179, 36, 34], + [180, 38, 34], + [181, 40, 33], + [182, 42, 33], + [183, 44, 33], + [184, 45, 33], + [185, 47, 32], + [186, 49, 32], + [187, 50, 32], + [188, 52, 31], + [189, 53, 31], + [190, 55, 31], + [191, 56, 31], + [192, 58, 30], + [193, 59, 30], + [194, 61, 30], + [195, 62, 29], + [196, 64, 29], + [197, 65, 28], + [198, 66, 28], + [199, 68, 28], + [200, 69, 27], + [201, 71, 27], + [202, 72, 26], + [203, 73, 26], + [204, 75, 25], + [205, 76, 25], + [206, 77, 24], + [207, 79, 24], + [208, 80, 23], + [209, 82, 23], + [210, 83, 22], + [211, 84, 21], + [212, 85, 21], + [213, 87, 20], + [214, 88, 19], + [215, 89, 19], + [216, 91, 18], + [217, 92, 17], + [218, 93, 16], + [219, 95, 15], + [220, 96, 14], + [221, 97, 13], + [222, 98, 12], + [223, 100, 11], + [224, 101, 9], + [225, 102, 8], + [226, 104, 7], + [227, 105, 5], + [227, 107, 5], + [227, 109, 6], + [228, 110, 7], + [228, 112, 7], + [228, 114, 8], + [228, 116, 8], + [229, 118, 9], + [229, 119, 10], + [229, 121, 10], + [229, 123, 11], + [229, 124, 12], + [230, 126, 12], + [230, 128, 13], + [230, 130, 14], + [230, 131, 14], + [230, 133, 15], + [230, 135, 15], + [231, 136, 16], + [231, 138, 17], + [231, 140, 17], + [231, 141, 18], + [231, 143, 19], + [231, 145, 19], + [231, 146, 20], + [232, 148, 21], + [232, 150, 21], + [232, 151, 22], + [232, 153, 22], + [232, 154, 23], + [232, 156, 24], + [232, 158, 24], + [232, 159, 25], + [232, 161, 26], + [232, 162, 26], + [233, 164, 27], + [233, 166, 27], + [233, 167, 28], + [233, 169, 29], + [233, 170, 29], + [233, 172, 30], + [233, 174, 30], + [233, 175, 31], + [233, 177, 32], + [233, 178, 32], + [233, 180, 33], + [233, 181, 34], + [233, 183, 34], + [233, 185, 35], + [233, 186, 35], + [233, 188, 36], + [233, 189, 37], + [233, 191, 37], + [233, 192, 38], + [233, 194, 38], + [233, 195, 39], + [233, 197, 40], + [233, 199, 40], + [233, 200, 41], + [232, 202, 42], + [232, 203, 42], + [232, 205, 43], + [232, 206, 43], + [232, 208, 44], + [232, 209, 45], + [232, 211, 45], + [232, 213, 46], + [232, 214, 47], + [232, 216, 47], + [231, 217, 48], + [231, 219, 48], + [231, 220, 49], + [231, 222, 50], + [231, 223, 50], + [231, 225, 51], + [230, 226, 52], + [230, 228, 52], + [230, 229, 53], + [231, 231, 60], + [233, 231, 69], + [234, 232, 78], + [236, 233, 87], + [237, 234, 94], + [238, 235, 102], + [240, 236, 109], + [241, 236, 117], + [242, 237, 124], + [243, 238, 131], + [245, 239, 137], + [246, 240, 144], + [247, 241, 151], + [248, 241, 158], + [249, 242, 164], + [249, 243, 171], + [250, 244, 177], + [251, 245, 184], + [252, 246, 190], + [252, 247, 197], + [253, 248, 203], + [253, 249, 210], + [254, 249, 216], + [254, 250, 223], + [254, 251, 229], + [255, 252, 236], + [255, 253, 242], + [255, 254, 249], + [255, 255, 255], +]; + +pub const EXTENDEDKINDLMANN: [[u8; 3]; 256] = [ + [0, 0, 0], + [5, 0, 4], + [9, 0, 9], + [13, 1, 13], + [16, 1, 17], + [19, 1, 21], + [22, 1, 24], + [24, 1, 27], + [26, 1, 30], + [27, 2, 34], + [28, 2, 38], + [29, 2, 42], + [29, 2, 46], + [30, 2, 50], + [30, 3, 53], + [30, 3, 57], + [30, 3, 61], + [29, 3, 65], + [29, 3, 68], + [28, 3, 72], + [27, 4, 75], + [27, 4, 79], + [26, 4, 82], + [25, 4, 85], + [24, 4, 88], + [23, 4, 92], + [22, 5, 95], + [21, 5, 98], + [20, 5, 101], + [19, 5, 103], + [18, 5, 106], + [18, 5, 109], + [17, 5, 111], + [14, 5, 115], + [8, 6, 119], + [6, 8, 120], + [6, 11, 120], + [6, 15, 119], + [6, 18, 118], + [6, 22, 116], + [5, 25, 114], + [5, 28, 112], + [5, 31, 109], + [5, 33, 107], + [5, 36, 104], + [5, 38, 101], + [5, 40, 99], + [5, 42, 96], + [5, 44, 94], + [4, 46, 91], + [4, 48, 89], + [4, 49, 87], + [4, 51, 85], + [4, 52, 83], + [4, 54, 81], + [4, 55, 79], + [4, 57, 77], + [4, 58, 76], + [4, 59, 74], + [3, 61, 73], + [3, 62, 71], + [3, 63, 70], + [3, 65, 69], + [3, 66, 67], + [3, 67, 66], + [3, 68, 65], + [3, 69, 64], + [3, 71, 63], + [3, 72, 61], + [4, 73, 60], + [4, 74, 58], + [4, 75, 56], + [4, 77, 55], + [4, 78, 53], + [4, 79, 51], + [4, 80, 49], + [4, 81, 47], + [4, 82, 45], + [4, 84, 43], + [4, 85, 41], + [4, 86, 39], + [4, 87, 36], + [4, 88, 34], + [4, 89, 31], + [4, 91, 29], + [4, 92, 26], + [5, 93, 24], + [5, 94, 21], + [5, 95, 18], + [5, 96, 15], + [5, 97, 13], + [5, 98, 10], + [5, 100, 8], + [5, 101, 6], + [6, 102, 5], + [8, 103, 5], + [10, 104, 5], + [11, 105, 5], + [13, 106, 5], + [15, 107, 5], + [17, 108, 5], + [20, 109, 5], + [22, 110, 5], + [26, 111, 5], + [29, 112, 5], + [32, 112, 5], + [36, 113, 5], + [40, 114, 5], + [43, 115, 6], + [47, 116, 6], + [51, 116, 6], + [55, 117, 6], + [59, 118, 6], + [63, 118, 6], + [67, 119, 6], + [71, 119, 6], + [76, 120, 6], + [80, 120, 6], + [84, 121, 6], + [88, 121, 6], + [92, 122, 6], + [97, 122, 6], + [101, 122, 6], + [105, 123, 6], + [109, 123, 6], + [113, 123, 6], + [118, 123, 6], + [122, 123, 6], + [126, 123, 6], + [130, 123, 6], + [135, 123, 7], + [139, 123, 7], + [144, 123, 7], + [149, 122, 7], + [154, 122, 7], + [160, 121, 8], + [165, 120, 8], + [171, 119, 8], + [177, 118, 8], + [183, 117, 9], + [189, 115, 9], + [196, 113, 9], + [202, 111, 10], + [209, 108, 10], + [216, 105, 10], + [222, 102, 11], + [229, 98, 11], + [236, 94, 11], + [243, 90, 12], + [244, 91, 27], + [245, 92, 37], + [245, 94, 46], + [245, 96, 52], + [246, 98, 58], + [246, 99, 63], + [246, 101, 67], + [246, 103, 71], + [246, 105, 74], + [246, 107, 77], + [247, 109, 79], + [247, 111, 83], + [247, 112, 87], + [247, 114, 91], + [247, 115, 96], + [248, 117, 101], + [248, 118, 106], + [248, 120, 112], + [248, 121, 118], + [249, 122, 123], + [249, 123, 129], + [249, 125, 135], + [249, 126, 141], + [249, 127, 147], + [249, 129, 153], + [249, 130, 158], + [249, 131, 164], + [249, 132, 169], + [249, 133, 175], + [249, 134, 180], + [249, 135, 186], + [249, 137, 191], + [249, 138, 196], + [249, 139, 201], + [250, 140, 206], + [250, 141, 211], + [250, 142, 216], + [250, 143, 221], + [250, 144, 225], + [250, 145, 230], + [250, 146, 235], + [250, 147, 239], + [250, 148, 243], + [250, 149, 248], + [248, 152, 250], + [245, 155, 250], + [242, 159, 250], + [239, 162, 251], + [236, 165, 251], + [234, 168, 251], + [232, 171, 251], + [230, 173, 251], + [229, 176, 251], + [227, 178, 251], + [226, 180, 251], + [225, 182, 251], + [224, 184, 252], + [224, 186, 252], + [223, 188, 252], + [223, 190, 252], + [222, 191, 252], + [222, 193, 252], + [222, 195, 252], + [222, 196, 252], + [223, 198, 252], + [223, 199, 252], + [223, 201, 252], + [224, 202, 252], + [224, 204, 253], + [225, 205, 253], + [226, 207, 253], + [226, 208, 253], + [227, 209, 253], + [228, 211, 253], + [229, 212, 253], + [229, 213, 253], + [230, 215, 253], + [231, 216, 253], + [231, 218, 253], + [231, 219, 253], + [232, 221, 253], + [232, 222, 253], + [232, 224, 253], + [232, 225, 254], + [233, 227, 254], + [233, 228, 254], + [233, 230, 254], + [233, 231, 254], + [233, 233, 254], + [233, 234, 254], + [234, 236, 254], + [234, 237, 254], + [234, 239, 254], + [235, 240, 254], + [235, 242, 254], + [236, 243, 254], + [237, 245, 254], + [237, 246, 254], + [238, 247, 254], + [239, 249, 254], + [240, 250, 254], + [242, 251, 254], + [243, 252, 254], + [245, 253, 255], + [248, 254, 255], + [251, 255, 255], + [255, 255, 255], +]; + +pub const KINDLMANN: [[u8; 3]; 256] = [ + [0, 0, 0], + [5, 0, 4], + [9, 0, 8], + [13, 1, 13], + [17, 1, 16], + [20, 1, 20], + [22, 1, 23], + [25, 1, 26], + [27, 1, 29], + [29, 2, 32], + [30, 2, 35], + [31, 2, 38], + [32, 2, 42], + [33, 2, 45], + [34, 2, 48], + [35, 2, 51], + [36, 3, 54], + [37, 3, 57], + [37, 3, 60], + [38, 3, 63], + [38, 3, 66], + [38, 3, 68], + [39, 3, 71], + [39, 4, 74], + [39, 4, 77], + [39, 4, 80], + [39, 4, 83], + [39, 4, 86], + [39, 4, 89], + [39, 4, 93], + [39, 5, 96], + [39, 5, 99], + [38, 5, 102], + [38, 5, 106], + [37, 5, 109], + [37, 5, 112], + [36, 6, 116], + [35, 6, 119], + [34, 6, 123], + [33, 6, 126], + [32, 6, 129], + [32, 6, 132], + [31, 6, 136], + [30, 7, 139], + [29, 7, 142], + [28, 7, 145], + [27, 7, 148], + [26, 7, 151], + [25, 7, 154], + [25, 7, 157], + [24, 8, 160], + [24, 8, 163], + [20, 8, 166], + [15, 8, 170], + [8, 9, 174], + [8, 12, 175], + [8, 15, 175], + [8, 19, 175], + [8, 22, 175], + [8, 25, 175], + [8, 28, 174], + [8, 32, 173], + [8, 35, 172], + [8, 38, 170], + [8, 41, 168], + [8, 43, 166], + [8, 46, 165], + [8, 49, 163], + [8, 51, 160], + [8, 53, 158], + [8, 56, 156], + [7, 58, 154], + [7, 60, 152], + [7, 62, 149], + [7, 64, 147], + [7, 66, 145], + [7, 68, 143], + [7, 70, 141], + [7, 71, 139], + [7, 73, 137], + [7, 75, 135], + [6, 76, 133], + [6, 78, 132], + [6, 80, 130], + [6, 81, 128], + [6, 83, 127], + [6, 84, 125], + [6, 86, 123], + [6, 87, 122], + [6, 88, 120], + [6, 90, 119], + [6, 91, 118], + [6, 93, 116], + [6, 94, 115], + [5, 95, 114], + [6, 96, 113], + [5, 98, 112], + [5, 99, 110], + [5, 100, 109], + [5, 102, 108], + [5, 103, 107], + [5, 104, 106], + [5, 105, 105], + [5, 107, 104], + [5, 108, 103], + [5, 109, 102], + [5, 110, 101], + [5, 112, 100], + [5, 113, 99], + [5, 114, 98], + [6, 115, 96], + [6, 117, 95], + [6, 118, 94], + [6, 119, 92], + [6, 120, 91], + [6, 122, 89], + [6, 123, 88], + [6, 124, 86], + [6, 125, 84], + [6, 127, 83], + [6, 128, 81], + [6, 129, 79], + [6, 130, 77], + [6, 132, 75], + [6, 133, 74], + [7, 134, 72], + [6, 135, 70], + [7, 137, 67], + [7, 138, 65], + [7, 139, 63], + [7, 140, 61], + [7, 142, 59], + [7, 143, 56], + [7, 144, 54], + [7, 145, 52], + [7, 147, 49], + [7, 148, 47], + [7, 149, 44], + [7, 150, 42], + [7, 151, 39], + [7, 153, 37], + [7, 154, 34], + [8, 155, 31], + [8, 156, 29], + [8, 157, 26], + [8, 159, 23], + [8, 160, 20], + [8, 161, 18], + [8, 162, 15], + [8, 163, 12], + [8, 165, 10], + [8, 166, 8], + [12, 167, 8], + [15, 168, 8], + [17, 169, 8], + [18, 170, 8], + [20, 171, 8], + [22, 173, 8], + [24, 174, 8], + [26, 175, 8], + [29, 176, 8], + [32, 177, 9], + [35, 178, 9], + [38, 179, 9], + [41, 180, 9], + [45, 181, 9], + [48, 182, 9], + [52, 183, 9], + [56, 184, 9], + [59, 185, 9], + [63, 186, 9], + [67, 187, 9], + [71, 188, 9], + [75, 189, 9], + [79, 190, 9], + [83, 190, 9], + [87, 191, 9], + [91, 192, 9], + [95, 193, 9], + [99, 193, 9], + [103, 194, 9], + [107, 195, 9], + [111, 196, 9], + [116, 196, 9], + [120, 197, 9], + [124, 198, 10], + [128, 198, 10], + [133, 199, 10], + [137, 199, 10], + [141, 200, 10], + [145, 200, 10], + [150, 201, 10], + [154, 201, 10], + [158, 202, 10], + [163, 202, 10], + [167, 202, 10], + [171, 203, 10], + [175, 203, 10], + [180, 203, 10], + [184, 204, 10], + [188, 204, 10], + [193, 204, 10], + [197, 205, 10], + [201, 205, 10], + [205, 205, 10], + [209, 205, 10], + [214, 205, 10], + [218, 205, 11], + [223, 205, 11], + [228, 205, 11], + [233, 205, 11], + [237, 205, 11], + [243, 205, 12], + [246, 204, 57], + [247, 205, 86], + [248, 205, 105], + [249, 206, 119], + [249, 207, 131], + [250, 207, 141], + [250, 208, 149], + [250, 209, 157], + [251, 210, 163], + [251, 211, 169], + [251, 212, 174], + [251, 214, 179], + [252, 215, 184], + [252, 216, 188], + [252, 217, 192], + [252, 218, 195], + [252, 220, 199], + [252, 221, 202], + [253, 222, 205], + [253, 224, 208], + [253, 225, 211], + [253, 226, 213], + [253, 227, 216], + [253, 229, 218], + [253, 230, 221], + [253, 232, 223], + [254, 233, 225], + [254, 234, 227], + [254, 236, 229], + [254, 237, 231], + [254, 238, 233], + [254, 240, 235], + [254, 241, 237], + [254, 242, 239], + [254, 244, 241], + [254, 245, 243], + [254, 247, 245], + [255, 248, 246], + [255, 249, 248], + [255, 251, 250], + [255, 252, 252], + [255, 254, 253], + [255, 255, 255], +]; + +pub const SMOOTHCOOLWARM: [[u8; 3]; 256] = [ + [59, 76, 192], + [60, 78, 194], + [61, 80, 195], + [62, 81, 197], + [63, 83, 199], + [64, 85, 200], + [66, 87, 202], + [67, 88, 203], + [68, 90, 204], + [69, 92, 206], + [70, 94, 207], + [72, 95, 209], + [73, 97, 210], + [74, 99, 212], + [75, 100, 213], + [76, 102, 214], + [78, 104, 216], + [79, 106, 217], + [80, 107, 218], + [81, 109, 219], + [83, 111, 221], + [84, 112, 222], + [85, 114, 223], + [86, 116, 224], + [88, 117, 226], + [89, 119, 227], + [90, 120, 228], + [91, 122, 229], + [93, 124, 230], + [94, 125, 231], + [95, 127, 232], + [97, 129, 233], + [98, 130, 234], + [99, 132, 235], + [101, 133, 236], + [102, 135, 237], + [103, 136, 238], + [105, 138, 239], + [106, 140, 240], + [107, 141, 240], + [109, 143, 241], + [110, 144, 242], + [111, 146, 243], + [113, 147, 244], + [114, 149, 244], + [115, 150, 245], + [117, 152, 246], + [118, 153, 246], + [119, 154, 247], + [121, 156, 248], + [122, 157, 248], + [123, 159, 249], + [125, 160, 249], + [126, 162, 250], + [128, 163, 250], + [129, 164, 251], + [130, 166, 251], + [132, 167, 252], + [133, 168, 252], + [134, 170, 252], + [136, 171, 253], + [137, 172, 253], + [139, 174, 253], + [140, 175, 254], + [141, 176, 254], + [143, 177, 254], + [144, 178, 254], + [146, 180, 254], + [147, 181, 255], + [148, 182, 255], + [150, 183, 255], + [151, 184, 255], + [153, 186, 255], + [154, 187, 255], + [155, 188, 255], + [157, 189, 255], + [158, 190, 255], + [159, 191, 255], + [161, 192, 255], + [162, 193, 255], + [163, 194, 254], + [165, 195, 254], + [166, 196, 254], + [168, 197, 254], + [169, 198, 254], + [170, 199, 253], + [172, 200, 253], + [173, 201, 253], + [174, 201, 252], + [176, 202, 252], + [177, 203, 252], + [178, 204, 251], + [180, 205, 251], + [181, 206, 250], + [182, 206, 250], + [183, 207, 249], + [185, 208, 249], + [186, 209, 248], + [187, 209, 248], + [189, 210, 247], + [190, 211, 246], + [191, 211, 246], + [192, 212, 245], + [193, 212, 244], + [195, 213, 244], + [196, 214, 243], + [197, 214, 242], + [198, 215, 241], + [200, 215, 241], + [201, 216, 240], + [202, 216, 239], + [203, 216, 238], + [204, 217, 237], + [205, 217, 236], + [206, 218, 235], + [208, 218, 234], + [209, 218, 233], + [210, 219, 232], + [211, 219, 231], + [212, 219, 230], + [213, 219, 229], + [214, 220, 228], + [215, 220, 227], + [216, 220, 226], + [217, 220, 225], + [218, 220, 224], + [219, 221, 222], + [220, 221, 221], + [221, 220, 220], + [222, 220, 219], + [223, 220, 217], + [225, 219, 216], + [226, 218, 214], + [227, 218, 213], + [228, 217, 211], + [229, 217, 210], + [229, 216, 209], + [230, 216, 207], + [231, 215, 206], + [232, 214, 204], + [233, 214, 203], + [234, 213, 201], + [235, 212, 200], + [235, 211, 198], + [236, 211, 197], + [237, 210, 195], + [238, 209, 194], + [238, 208, 192], + [239, 207, 191], + [239, 206, 189], + [240, 206, 187], + [241, 205, 186], + [241, 204, 184], + [242, 203, 183], + [242, 202, 181], + [243, 201, 180], + [243, 200, 178], + [244, 199, 177], + [244, 198, 175], + [244, 197, 173], + [245, 196, 172], + [245, 195, 170], + [245, 193, 169], + [246, 192, 167], + [246, 191, 166], + [246, 190, 164], + [246, 189, 162], + [247, 188, 161], + [247, 186, 159], + [247, 185, 158], + [247, 184, 156], + [247, 183, 155], + [247, 181, 153], + [247, 180, 151], + [247, 179, 150], + [247, 177, 148], + [247, 176, 147], + [247, 175, 145], + [247, 173, 144], + [247, 172, 142], + [247, 171, 140], + [247, 169, 139], + [247, 168, 137], + [247, 166, 136], + [246, 165, 134], + [246, 163, 133], + [246, 162, 131], + [246, 160, 129], + [245, 159, 128], + [245, 157, 126], + [245, 156, 125], + [244, 154, 123], + [244, 153, 122], + [244, 151, 120], + [243, 149, 119], + [243, 148, 117], + [242, 146, 116], + [242, 144, 114], + [241, 143, 113], + [241, 141, 111], + [240, 139, 110], + [240, 138, 108], + [239, 136, 107], + [238, 134, 105], + [238, 133, 104], + [237, 131, 102], + [237, 129, 101], + [236, 127, 99], + [235, 125, 98], + [234, 124, 96], + [234, 122, 95], + [233, 120, 93], + [232, 118, 92], + [231, 116, 91], + [230, 114, 89], + [229, 112, 88], + [229, 111, 86], + [228, 109, 85], + [227, 107, 84], + [226, 105, 82], + [225, 103, 81], + [224, 101, 79], + [223, 99, 78], + [222, 97, 77], + [221, 95, 75], + [220, 93, 74], + [219, 91, 73], + [218, 89, 71], + [216, 86, 70], + [215, 84, 69], + [214, 82, 68], + [213, 80, 66], + [212, 78, 65], + [211, 76, 64], + [209, 73, 62], + [208, 71, 61], + [207, 69, 60], + [206, 67, 59], + [204, 64, 57], + [203, 62, 56], + [202, 59, 55], + [200, 57, 54], + [199, 54, 53], + [198, 52, 52], + [196, 49, 50], + [195, 46, 49], + [193, 43, 48], + [192, 40, 47], + [191, 37, 46], + [189, 34, 45], + [188, 30, 44], + [186, 26, 43], + [185, 22, 41], + [183, 17, 40], + [182, 11, 39], + [180, 4, 38], +]; diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 3d418d9..893343a 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,34 +1,32 @@ -use anyhow::Result; +use anyhow::{anyhow, Result}; use indicatif::{ProgressBar, ProgressStyle}; use std::io::{Read, Write}; use std::path::{Path, PathBuf}; pub mod coco; -mod turbo; +pub mod colormap256; -pub use turbo::TURBO; +pub use colormap256::*; pub const GITHUB_ASSETS: &str = "https://github.com/jamjamjon/assets/releases/download/v0.0.1"; pub const CHECK_MARK: &str = "✅"; pub const CROSS_MARK: &str = "❌"; pub const SAFE_CROSS_MARK: &str = "❎"; -pub fn auto_load>(src: P) -> Result { +pub fn auto_load>(src: P, sub: Option<&str>) -> Result { let src = src.as_ref(); let p = if src.is_file() { src.into() } else { let sth = src.file_name().unwrap().to_str().unwrap(); - let mut p = config_dir(); + let mut p = home_dir(sub); p.push(sth); - // download from github assets if not exists in config directory if !p.is_file() { download( &format!("{}/{}", GITHUB_ASSETS, sth), &p, Some(sth.to_string().as_str()), - ) - .unwrap_or_else(|err| panic!("Fail to load {:?}: {err}", src)); + )?; } p }; @@ -46,7 +44,7 @@ pub fn download + std::fmt::Debug>( .get(src) .timeout(std::time::Duration::from_secs(2000)) .call() - .unwrap_or_else(|err| panic!("Failed to GET: {}", err)); + .map_err(|err| anyhow!("Failed to download. {err:?}"))?; let ntotal = resp .header("Content-Length") .and_then(|s| s.parse::().ok()) @@ -93,7 +91,23 @@ pub fn config_dir() -> PathBuf { Some(mut d) => { d.push("usls"); if !d.exists() { - std::fs::create_dir_all(&d).expect("Failed to create config directory."); + std::fs::create_dir_all(&d).expect("Failed to create usls config directory."); + } + d + } + None => panic!("Unsupported operating system. Now support Linux, MacOS, Windows."), + } +} + +pub fn home_dir(sub: Option<&str>) -> PathBuf { + match dirs::home_dir() { + Some(mut d) => { + d.push(".usls"); + if let Some(sub) = sub { + d.push(sub); + } + if !d.exists() { + std::fs::create_dir_all(&d).expect("Failed to create usls home directory."); } d } diff --git a/src/utils/turbo.rs b/src/utils/turbo.rs deleted file mode 100644 index e9e5518..0000000 --- a/src/utils/turbo.rs +++ /dev/null @@ -1,258 +0,0 @@ -pub const TURBO: [[u8; 3]; 256] = [ - [48, 18, 59], - [50, 21, 67], - [51, 24, 74], - [52, 27, 81], - [53, 30, 88], - [54, 33, 95], - [55, 36, 102], - [56, 39, 109], - [57, 42, 115], - [58, 45, 121], - [59, 47, 128], - [60, 50, 134], - [61, 53, 139], - [62, 56, 145], - [63, 59, 151], - [63, 62, 156], - [64, 64, 162], - [65, 67, 167], - [65, 70, 172], - [66, 73, 177], - [66, 75, 181], - [67, 78, 186], - [68, 81, 191], - [68, 84, 195], - [68, 86, 199], - [69, 89, 203], - [69, 92, 207], - [69, 94, 211], - [70, 97, 214], - [70, 100, 218], - [70, 102, 221], - [70, 105, 224], - [70, 107, 227], - [71, 110, 230], - [71, 113, 233], - [71, 115, 235], - [71, 118, 238], - [71, 120, 240], - [71, 123, 242], - [70, 125, 244], - [70, 128, 246], - [70, 130, 248], - [70, 133, 250], - [70, 135, 251], - [69, 138, 252], - [69, 140, 253], - [68, 143, 254], - [67, 145, 254], - [66, 148, 255], - [65, 150, 255], - [64, 153, 255], - [62, 155, 254], - [61, 158, 254], - [59, 160, 253], - [58, 163, 252], - [56, 165, 251], - [55, 168, 250], - [53, 171, 248], - [51, 173, 247], - [49, 175, 245], - [47, 178, 244], - [46, 180, 242], - [44, 183, 240], - [42, 185, 238], - [40, 188, 235], - [39, 190, 233], - [37, 192, 231], - [35, 195, 228], - [34, 197, 226], - [32, 199, 223], - [31, 201, 221], - [30, 203, 218], - [28, 205, 216], - [27, 208, 213], - [26, 210, 210], - [26, 212, 208], - [25, 213, 205], - [24, 215, 202], - [24, 217, 200], - [24, 219, 197], - [24, 221, 194], - [24, 222, 192], - [24, 224, 189], - [25, 226, 187], - [25, 227, 185], - [26, 228, 182], - [28, 230, 180], - [29, 231, 178], - [31, 233, 175], - [32, 234, 172], - [34, 235, 170], - [37, 236, 167], - [39, 238, 164], - [42, 239, 161], - [44, 240, 158], - [47, 241, 155], - [50, 242, 152], - [53, 243, 148], - [56, 244, 145], - [60, 245, 142], - [63, 246, 138], - [67, 247, 135], - [70, 248, 132], - [74, 248, 128], - [78, 249, 125], - [82, 250, 122], - [85, 250, 118], - [89, 251, 115], - [93, 252, 111], - [97, 252, 108], - [101, 253, 105], - [105, 253, 102], - [109, 254, 98], - [113, 254, 95], - [117, 254, 92], - [121, 254, 89], - [125, 255, 86], - [128, 255, 83], - [132, 255, 81], - [136, 255, 78], - [139, 255, 75], - [143, 255, 73], - [146, 255, 71], - [150, 254, 68], - [153, 254, 66], - [156, 254, 64], - [159, 253, 63], - [161, 253, 61], - [164, 252, 60], - [167, 252, 58], - [169, 251, 57], - [172, 251, 56], - [175, 250, 55], - [177, 249, 54], - [180, 248, 54], - [183, 247, 53], - [185, 246, 53], - [188, 245, 52], - [190, 244, 52], - [193, 243, 52], - [195, 241, 52], - [198, 240, 52], - [200, 239, 52], - [203, 237, 52], - [205, 236, 52], - [208, 234, 52], - [210, 233, 53], - [212, 231, 53], - [215, 229, 53], - [217, 228, 54], - [219, 226, 54], - [221, 224, 55], - [223, 223, 55], - [225, 221, 55], - [227, 219, 56], - [229, 217, 56], - [231, 215, 57], - [233, 213, 57], - [235, 211, 57], - [236, 209, 58], - [238, 207, 58], - [239, 205, 58], - [241, 203, 58], - [242, 201, 58], - [244, 199, 58], - [245, 197, 58], - [246, 195, 58], - [247, 193, 58], - [248, 190, 57], - [249, 188, 57], - [250, 186, 57], - [251, 184, 56], - [251, 182, 55], - [252, 179, 54], - [252, 177, 54], - [253, 174, 53], - [253, 172, 52], - [254, 169, 51], - [254, 167, 50], - [254, 164, 49], - [254, 161, 48], - [254, 158, 47], - [254, 155, 45], - [254, 153, 44], - [254, 150, 43], - [254, 147, 42], - [254, 144, 41], - [253, 141, 39], - [253, 138, 38], - [252, 135, 37], - [252, 132, 35], - [251, 129, 34], - [251, 126, 33], - [250, 123, 31], - [249, 120, 30], - [249, 117, 29], - [248, 114, 28], - [247, 111, 26], - [246, 108, 25], - [245, 105, 24], - [244, 102, 23], - [243, 99, 21], - [242, 96, 20], - [241, 93, 19], - [240, 91, 18], - [239, 88, 17], - [237, 85, 16], - [236, 83, 15], - [235, 80, 14], - [234, 78, 13], - [232, 75, 12], - [231, 73, 12], - [229, 71, 11], - [228, 69, 10], - [226, 67, 10], - [225, 65, 9], - [223, 63, 8], - [221, 61, 8], - [220, 59, 7], - [218, 57, 7], - [216, 55, 6], - [214, 53, 6], - [212, 51, 5], - [210, 49, 5], - [208, 47, 5], - [206, 45, 4], - [204, 43, 4], - [202, 42, 4], - [200, 40, 3], - [197, 38, 3], - [195, 37, 3], - [193, 35, 2], - [190, 33, 2], - [188, 32, 2], - [185, 30, 2], - [183, 29, 2], - [180, 27, 1], - [178, 26, 1], - [175, 24, 1], - [172, 23, 1], - [169, 22, 1], - [167, 20, 1], - [164, 19, 1], - [161, 18, 1], - [158, 16, 1], - [155, 15, 1], - [152, 14, 1], - [149, 13, 1], - [146, 11, 1], - [142, 10, 1], - [139, 9, 2], - [136, 8, 2], - [133, 7, 2], - [129, 6, 2], - [126, 5, 2], - [122, 4, 3], -]; diff --git a/src/ys/bbox.rs b/src/ys/bbox.rs index 3612299..856cdc3 100644 --- a/src/ys/bbox.rs +++ b/src/ys/bbox.rs @@ -161,6 +161,27 @@ impl Bbox { self.confidence } + pub fn label(&self, with_name: bool, with_conf: bool) -> String { + let mut label = String::new(); + if with_name { + label.push_str( + &self + .name + .as_ref() + .unwrap_or(&self.id.to_string()) + .to_string(), + ); + } + if with_conf { + if with_name { + label.push_str(&format!(": {:.4}", self.confidence)); + } else { + label.push_str(&format!("{:.4}", self.confidence)); + } + } + label + } + pub fn area(&self) -> f32 { self.h * self.w } diff --git a/src/ys/keypoint.rs b/src/ys/keypoint.rs index 8d3fc9c..9476882 100644 --- a/src/ys/keypoint.rs +++ b/src/ys/keypoint.rs @@ -215,6 +215,27 @@ impl Keypoint { self.name.as_ref() } + pub fn label(&self, with_name: bool, with_conf: bool) -> String { + let mut label = String::new(); + if with_name { + label.push_str( + &self + .name + .as_ref() + .unwrap_or(&self.id.to_string()) + .to_string(), + ); + } + if with_conf { + if with_name { + label.push_str(&format!(": {:.4}", self.confidence)); + } else { + label.push_str(&format!("{:.4}", self.confidence)); + } + } + label + } + pub fn is_origin(&self) -> bool { self.x == 0.0_f32 && self.y == 0.0_f32 } diff --git a/src/ys/mask.rs b/src/ys/mask.rs index 04691a6..89c6476 100644 --- a/src/ys/mask.rs +++ b/src/ys/mask.rs @@ -1,23 +1,19 @@ -use geo::{ - coord, point, polygon, Area, BoundingRect, Centroid, ConvexHull, EuclideanLength, LineString, - MinimumRotatedRect, Point, Polygon, Simplify, -}; +use image::DynamicImage; -use crate::{Bbox, Mbr}; - -/// Mask #[derive(Clone, PartialEq)] pub struct Mask { - polygon: Polygon, + mask: DynamicImage, + mask_vec: Vec, id: isize, name: Option, - confidence: f32, + confidence: f32, // placeholder } impl Default for Mask { fn default() -> Self { Self { - polygon: polygon![], + mask: DynamicImage::default(), + mask_vec: vec![], id: -1, name: None, confidence: 0., @@ -28,27 +24,22 @@ impl Default for Mask { impl std::fmt::Debug for Mask { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Mask") - // .field("polygons", &self.polygon) + // .field("mask", &self.mask) .field("id", &self.id) .field("name", &self.name) - .field("confidence", &self.confidence) + // .field("confidence", &self.confidence) .finish() } } impl Mask { - pub fn with_points_imageproc(mut self, points: &[imageproc::point::Point]) -> Self { - // exterior - let v = points - .iter() - .map(|p| coord! { x: p.x as f64, y: p.y as f64}) - .collect::>(); - self.polygon = Polygon::new(LineString::from(v), vec![]); + pub fn with_mask(mut self, x: DynamicImage) -> Self { + self.mask = x; self } - pub fn with_polygon(mut self, x: Polygon) -> Self { - self.polygon = x; + pub fn with_vec(mut self, vec: &[u8]) -> Self { + self.mask_vec = vec.to_vec(); self } @@ -62,6 +53,15 @@ impl Mask { self } + pub fn mask(&self) -> &DynamicImage { + &self.mask + } + + pub fn vec(&self) -> Vec { + // self.mask.to_luma8().into_raw() + self.mask_vec.clone() + } + pub fn id(&self) -> isize { self.id } @@ -73,117 +73,4 @@ impl Mask { pub fn confidence(&self) -> f32 { self.confidence } - - pub fn polygon(&self) -> &Polygon { - &self.polygon - } - - pub fn is_closed(&self) -> bool { - self.polygon.exterior().is_closed() - } - - pub fn count(&self) -> usize { - self.polygon.exterior().points().len() - } - - pub fn perimeter(&self) -> f64 { - self.polygon.exterior().euclidean_length() - } - - pub fn area(&self) -> f64 { - self.polygon.unsigned_area() - } - - pub fn centroid(&self) -> Option<(f32, f32)> { - self.polygon - .centroid() - .map(|x| (x.x() as f32, x.y() as f32)) - } - - pub fn bbox(&self) -> Option { - self.polygon.bounding_rect().map(|x| { - Bbox::default().with_xyxy( - x.min().x as f32, - x.min().y as f32, - x.max().x as f32, - x.max().y as f32, - ) - }) - } - - pub fn mbr(&self) -> Option { - MinimumRotatedRect::minimum_rotated_rect(&self.polygon) - .map(|x| Mbr::from_line_string(x.exterior().to_owned())) - } - - pub fn convex_hull(mut self) -> Self { - self.polygon = self.polygon.convex_hull(); - self - } - - pub fn simplify(mut self, eps: f64) -> Self { - self.polygon = self.polygon.simplify(&eps); - self - } - - pub fn resample(mut self, num_samples: usize) -> Self { - let points = self.polygon.exterior().to_owned().into_points(); - let mut new_points = Vec::new(); - for i in 0..points.len() { - let start_point: Point = points[i]; - let end_point = points[(i + 1) % points.len()]; - new_points.push(start_point); - let dx = end_point.x() - start_point.x(); - let dy = end_point.y() - start_point.y(); - for j in 1..num_samples { - let t = (j as f64) / (num_samples as f64); - let new_x = start_point.x() + t * dx; - let new_y = start_point.y() + t * dy; - new_points.push(point! { x: new_x, y: new_y }); - } - } - self.polygon = Polygon::new(LineString::from(new_points), vec![]); - self - } - - pub fn unclip(mut self, delta: f64, width: f64, height: f64) -> Self { - let points = self.polygon.exterior().to_owned().into_points(); - let num_points = points.len(); - let mut new_points = Vec::with_capacity(points.len()); - for i in 0..num_points { - let prev_idx = if i == 0 { num_points - 1 } else { i - 1 }; - let next_idx = (i + 1) % num_points; - - let edge_vector = point! { - x: points[next_idx].x() - points[prev_idx].x(), - y: points[next_idx].y() - points[prev_idx].y(), - }; - let normal_vector = point! { - x: -edge_vector.y(), - y: edge_vector.x(), - }; - - let normal_length = (normal_vector.x().powi(2) + normal_vector.y().powi(2)).sqrt(); - if normal_length.abs() < 1e-6 { - new_points.push(points[i]); - } else { - let normalized_normal = point! { - x: normal_vector.x() / normal_length, - y: normal_vector.y() / normal_length, - }; - - let new_x = points[i].x() + normalized_normal.x() * delta; - let new_y = points[i].y() + normalized_normal.y() * delta; - let new_x = new_x.max(0.0).min(width); - let new_y = new_y.max(0.0).min(height); - new_points.push(point! { - x: new_x, - y: new_y, - }); - } - } - - self.polygon = Polygon::new(LineString::from(new_points), vec![]); - self - } } diff --git a/src/ys/mbr.rs b/src/ys/mbr.rs index 3417b3e..29cee23 100644 --- a/src/ys/mbr.rs +++ b/src/ys/mbr.rs @@ -104,6 +104,27 @@ impl Mbr { self.confidence } + pub fn label(&self, with_name: bool, with_conf: bool) -> String { + let mut label = String::new(); + if with_name { + label.push_str( + &self + .name + .as_ref() + .unwrap_or(&self.id.to_string()) + .to_string(), + ); + } + if with_conf { + if with_name { + label.push_str(&format!(": {:.4}", self.confidence)); + } else { + label.push_str(&format!("{:.4}", self.confidence)); + } + } + label + } + pub fn vertices(&self) -> Vec { self.ls.0.clone() } diff --git a/src/ys/mod.rs b/src/ys/mod.rs index d205719..010032c 100644 --- a/src/ys/mod.rs +++ b/src/ys/mod.rs @@ -3,6 +3,7 @@ mod embedding; mod keypoint; mod mask; mod mbr; +mod polygon; mod prob; mod y; @@ -11,5 +12,6 @@ pub use embedding::Embedding; pub use keypoint::Keypoint; pub use mask::Mask; pub use mbr::Mbr; +pub use polygon::Polygon; pub use prob::Prob; pub use y::Y; diff --git a/src/ys/polygon.rs b/src/ys/polygon.rs new file mode 100644 index 0000000..cfebd03 --- /dev/null +++ b/src/ys/polygon.rs @@ -0,0 +1,209 @@ +use geo::{ + coord, point, polygon, Area, BoundingRect, Centroid, ConvexHull, EuclideanLength, LineString, + MinimumRotatedRect, Point, Simplify, +}; + +use crate::{Bbox, Mbr}; + +#[derive(Clone, PartialEq)] +pub struct Polygon { + polygon: geo::Polygon, + id: isize, + name: Option, + confidence: f32, +} + +impl Default for Polygon { + fn default() -> Self { + Self { + polygon: polygon![], + id: -1, + name: None, + confidence: 0., + } + } +} + +impl std::fmt::Debug for Polygon { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Polygon") + // .field("polygons", &self.polygon) + .field("id", &self.id) + .field("name", &self.name) + .field("confidence", &self.confidence) + .finish() + } +} + +impl Polygon { + pub fn with_points_imageproc(mut self, points: &[imageproc::point::Point]) -> Self { + // exterior + let v = points + .iter() + .map(|p| coord! { x: p.x as f64, y: p.y as f64}) + .collect::>(); + self.polygon = geo::Polygon::new(LineString::from(v), vec![]); + self + } + + pub fn with_polygon(mut self, x: geo::Polygon) -> Self { + self.polygon = x; + self + } + + pub fn with_id(mut self, x: isize) -> Self { + self.id = x; + self + } + + pub fn with_name(mut self, x: Option) -> Self { + self.name = x; + self + } + + pub fn id(&self) -> isize { + self.id + } + + pub fn name(&self) -> Option<&String> { + self.name.as_ref() + } + + pub fn confidence(&self) -> f32 { + self.confidence + } + + pub fn label(&self, with_name: bool, with_conf: bool) -> String { + let mut label = String::new(); + if with_name { + label.push_str( + &self + .name + .as_ref() + .unwrap_or(&self.id.to_string()) + .to_string(), + ); + } + if with_conf { + if with_name { + label.push_str(&format!(": {:.4}", self.confidence)); + } else { + label.push_str(&format!("{:.4}", self.confidence)); + } + } + label + } + + pub fn polygon(&self) -> &geo::Polygon { + &self.polygon + } + + pub fn is_closed(&self) -> bool { + self.polygon.exterior().is_closed() + } + + pub fn count(&self) -> usize { + self.polygon.exterior().points().len() + } + + pub fn perimeter(&self) -> f64 { + self.polygon.exterior().euclidean_length() + } + + pub fn area(&self) -> f64 { + self.polygon.unsigned_area() + } + + pub fn centroid(&self) -> Option<(f32, f32)> { + self.polygon + .centroid() + .map(|x| (x.x() as f32, x.y() as f32)) + } + + pub fn bbox(&self) -> Option { + self.polygon.bounding_rect().map(|x| { + Bbox::default().with_xyxy( + x.min().x as f32, + x.min().y as f32, + x.max().x as f32, + x.max().y as f32, + ) + }) + } + + pub fn mbr(&self) -> Option { + MinimumRotatedRect::minimum_rotated_rect(&self.polygon) + .map(|x| Mbr::from_line_string(x.exterior().to_owned())) + } + + pub fn convex_hull(mut self) -> Self { + self.polygon = self.polygon.convex_hull(); + self + } + + pub fn simplify(mut self, eps: f64) -> Self { + self.polygon = self.polygon.simplify(&eps); + self + } + + pub fn resample(mut self, num_samples: usize) -> Self { + let points = self.polygon.exterior().to_owned().into_points(); + let mut new_points = Vec::new(); + for i in 0..points.len() { + let start_point: Point = points[i]; + let end_point = points[(i + 1) % points.len()]; + new_points.push(start_point); + let dx = end_point.x() - start_point.x(); + let dy = end_point.y() - start_point.y(); + for j in 1..num_samples { + let t = (j as f64) / (num_samples as f64); + let new_x = start_point.x() + t * dx; + let new_y = start_point.y() + t * dy; + new_points.push(point! { x: new_x, y: new_y }); + } + } + self.polygon = geo::Polygon::new(LineString::from(new_points), vec![]); + self + } + + pub fn unclip(mut self, delta: f64, width: f64, height: f64) -> Self { + let points = self.polygon.exterior().to_owned().into_points(); + let num_points = points.len(); + let mut new_points = Vec::with_capacity(points.len()); + for i in 0..num_points { + let prev_idx = if i == 0 { num_points - 1 } else { i - 1 }; + let next_idx = (i + 1) % num_points; + + let edge_vector = point! { + x: points[next_idx].x() - points[prev_idx].x(), + y: points[next_idx].y() - points[prev_idx].y(), + }; + let normal_vector = point! { + x: -edge_vector.y(), + y: edge_vector.x(), + }; + + let normal_length = (normal_vector.x().powi(2) + normal_vector.y().powi(2)).sqrt(); + if normal_length.abs() < 1e-6 { + new_points.push(points[i]); + } else { + let normalized_normal = point! { + x: normal_vector.x() / normal_length, + y: normal_vector.y() / normal_length, + }; + + let new_x = points[i].x() + normalized_normal.x() * delta; + let new_y = points[i].y() + normalized_normal.y() * delta; + let new_x = new_x.max(0.0).min(width); + let new_y = new_y.max(0.0).min(height); + new_points.push(point! { + x: new_x, + y: new_y, + }); + } + } + + self.polygon = geo::Polygon::new(LineString::from(new_points), vec![]); + self + } +} diff --git a/src/ys/y.rs b/src/ys/y.rs index ed1083d..24c5ec3 100644 --- a/src/ys/y.rs +++ b/src/ys/y.rs @@ -1,4 +1,4 @@ -use crate::{Bbox, Keypoint, Mask, Mbr, Prob}; +use crate::{Bbox, Keypoint, Mask, Mbr, Polygon, Prob}; #[derive(Clone, PartialEq, Default)] pub struct Y { @@ -6,14 +6,14 @@ pub struct Y { bboxes: Option>, keypoints: Option>>, mbrs: Option>, - masks: Option>, + polygons: Option>, texts: Option>, - pixels: Option>, + masks: Option>, } impl std::fmt::Debug for Y { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let mut f = f.debug_struct("Results"); + let mut f = f.debug_struct("Result"); if let Some(x) = &self.texts { if !x.is_empty() { f.field("Texts", &x); @@ -37,23 +37,23 @@ impl std::fmt::Debug for Y { f.field("Keypoints", &x); } } + if let Some(x) = &self.polygons { + if !x.is_empty() { + f.field("Polygons", &x); + } + } if let Some(x) = &self.masks { if !x.is_empty() { f.field("Masks", &x); } } - if let Some(x) = &self.pixels { - if !x.is_empty() { - f.field("Pixels", &x); - } - } f.finish() } } impl Y { - pub fn with_pixels(mut self, pixels: &[u8]) -> Self { - self.pixels = Some(pixels.to_vec()); + pub fn with_masks(mut self, masks: &[Mask]) -> Self { + self.masks = Some(masks.to_vec()); self } @@ -81,13 +81,13 @@ impl Y { self } - pub fn with_masks(mut self, masks: &[Mask]) -> Self { - self.masks = Some(masks.to_vec()); + pub fn with_polygons(mut self, polygons: &[Polygon]) -> Self { + self.polygons = Some(polygons.to_vec()); self } - pub fn pixels(&self) -> Option<&Vec> { - self.pixels.as_ref() + pub fn masks(&self) -> Option<&Vec> { + self.masks.as_ref() } pub fn probs(&self) -> Option<&Prob> { @@ -98,8 +98,8 @@ impl Y { self.keypoints.as_ref() } - pub fn masks(&self) -> Option<&Vec> { - self.masks.as_ref() + pub fn polygons(&self) -> Option<&Vec> { + self.polygons.as_ref() } pub fn bboxes(&self) -> Option<&Vec> {