Add MODNet model (#11)

* Add MODNet for portrait matting

* Minor fixes

* Move assets to home directory

* Add colormap

* ci

* Update README.md

.github/workflows/rust-ci.yml

@@ -2,7 +2,7 @@ name: Rust
 
 on:
   push:
-    branches: [ "main" ]
+    branches: [ "main", "dev" ]
   pull_request:
     branches: [ "main" ]
 

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"

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         |
+| :----------------------------: |
+|<img src='examples/modnet/demo.png'   width="800px">|
+
+
 |        Depth-Anything         |
 | :----------------------------: |
 |<img src='examples/depth-anything/demo.png'   width="800px">|
 
 
-
 |        YOLOP-v2          |             Face-Parsing             |               Text-Detection           |  
 | :----------------------------: | :------------------------------: |  :------------------------------: |
-|<img src='examples/yolop/demo.png'  height="240px">| <img src='examples/face-parsing/demo.png'  height="240px"> | <img src='examples/db/demo.png'  height="240px"> |
+|<img src='examples/yolop/demo.png'  height="180px">| <img src='examples/face-parsing/demo.png'  height="180px"> | <img src='examples/db/demo.png'  height="180px"> |
 
 
 |        YOLOv8-Obb         |
@@ -23,8 +27,6 @@ A Rust library integrated with **ONNXRuntime**, providing a collection of **Comp
 
 
 
-
-
 ## Supported Models
 
 |                               Model                               |         Task / Type         |         Example         | CUDA<br />f32 | CUDA<br />f16 |     TensorRT<br />f32     |     TensorRT<br />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
 

examples/blip/main.rs

@@ -3,14 +3,14 @@ use usls::{models::Blip, DataLoader, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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

examples/clip/main.rs

@@ -3,14 +3,14 @@ use usls::{models::Clip, DataLoader, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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);
 

examples/db/main.rs

@@ -10,7 +10,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         .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<dyn std::error::Error>> {
     // 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);

examples/depth-anything/main.rs

@@ -3,7 +3,7 @@ use usls::{models::DepthAnything, Annotator, DataLoader, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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<dyn std::error::Error>> {
     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(())

examples/dinov2/main.rs

@@ -3,7 +3,7 @@ use usls::{models::Dinov2, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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());

examples/face-parsing/main.rs

@@ -3,7 +3,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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<dyn std::error::Error>> {
         .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);
 

examples/fastsam/main.rs

@@ -3,7 +3,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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())

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)

examples/modnet/main.rs

@@ -0,0 +1,23 @@
+use usls::{models::MODNet, Annotator, DataLoader, Options};
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    // 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(())
+}

examples/rtdetr/main.rs

@@ -3,7 +3,7 @@ use usls::{coco, models::RTDETR, Annotator, DataLoader, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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)?;

examples/rtmo/main.rs

@@ -3,7 +3,7 @@ use usls::{coco, models::RTMO, Annotator, DataLoader, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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])

examples/svtr/main.rs

@@ -6,8 +6,8 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         .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

examples/yolo-world/main.rs

@@ -3,7 +3,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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())

examples/yolop/main.rs

@@ -3,7 +3,7 @@ use usls::{models::YOLOPv2, Annotator, DataLoader, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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<dyn std::error::Error>> {
 
     // annotate
     let annotator = Annotator::default()
-        .with_masks_name(true)
+        .with_polygons_name(true)
         .with_saveout("YOLOPv2");
     annotator.annotate(&x, &y);
 

examples/yolov5/main.rs

@@ -9,7 +9,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         .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())

examples/yolov8-face/main.rs

@@ -3,7 +3,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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())

examples/yolov8-falldown/main.rs

@@ -2,7 +2,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options};
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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

examples/yolov8-head/main.rs

@@ -2,7 +2,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options};
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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

examples/yolov8-trash/main.rs

@@ -3,7 +3,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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)?;
 

examples/yolov8/main.rs

@@ -3,23 +3,20 @@ use usls::{coco, models::YOLO, Annotator, DataLoader, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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<dyn std::error::Error>> {
 
     // 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

examples/yolov9/main.rs

@@ -3,7 +3,7 @@ use usls::{models::YOLO, Annotator, DataLoader, Options};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 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())

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<String>,
+
     // About mbrs
     without_mbrs: bool,
     without_mbrs_conf: bool,
     without_mbrs_name: bool,
     without_mbrs_text_bg: bool,
     mbrs_text_color: Rgba<u8>,
+
     // About bboxes
     without_bboxes: bool,
     without_bboxes_conf: bool,
     without_bboxes_name: bool,
     without_bboxes_text_bg: bool,
     bboxes_text_color: Rgba<u8>,
+
     // About keypoints
     without_keypoints: bool,
     with_keypoints_conf: bool,
@@ -34,15 +38,21 @@ pub struct Annotator {
     skeletons: Option<Vec<(usize, usize)>>,
     keypoints_radius: usize,
     keypoints_palette: Option<Vec<(u8, u8, u8, u8)>>,
+
+    // 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<u8>,
+    polygons_alpha: u8,
+    contours_color: Rgba<u8>,
+
     // 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<u8>,
-    masks_alpha: u8,
-    polygon_color: Rgba<u8>,
+    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<image::Luma<_>, Vec<u8>> =
-            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::<Vec<_>>();
-            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::<Vec<_>>();
-                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::<Vec<_>>();
+            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::<Vec<_>>();
+                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<Keypoint>]) {
+        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<Keypoint>]) {
-        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<FontVec> {
         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

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)

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> {
+        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> {
+        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> {
+        self.tokenizer = Some(auto_load(tokenizer, Some("models"))?);
+        Ok(self)
     }
 
     pub fn with_unclip_ratio(mut self, x: f32) -> Self {

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<Mask> = Vec::new();
+            let mut y_polygons: Vec<Polygon> = Vec::new();
             let mut y_mbrs: Vec<Mbr> = 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),
             );
         }

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)
     }

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;

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<Self> {
+        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<Vec<Y>> {
+        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<Array<f32, IxDyn>>, xs0: &[DynamicImage]) -> Result<Vec<Y>> {
+        let mut ys: Vec<Y> = 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::<Vec<_>>();
+            let luma: image::ImageBuffer<image::Luma<_>, Vec<u8>> =
+                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
+    }
+}

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<Mask> = Vec::new();
+                            let mut y_polygons: Vec<Polygon> = 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<Mask> = Vec::new();
+                                let mut masks: Vec<Polygon> = Vec::new();
                                 let contours: Vec<imageproc::contours::Contour<i32>> =
                                     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);

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<Mask> = Vec::new();
+            let mut y_polygons: Vec<Polygon> = Vec::new();
             let contours: Vec<imageproc::contours::Contour<i32>> =
                 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::Contour<i32>> =
                 imageproc::contours::find_contours_with_threshold(&mask_ll, 1);
-            let mut masks: Vec<Mask> = Vec::new();
+            let mut masks: Vec<Polygon> = 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),
             );
         }

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<P: AsRef<Path>>(src: P) -> Result<String> {
+pub fn auto_load<P: AsRef<Path>>(src: P, sub: Option<&str>) -> Result<String> {
     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<P: AsRef<Path> + 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::<u64>().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
         }

src/utils/turbo.rs

@@ -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],
-];

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
     }

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
     }

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<u8>,
     id: isize,
     name: Option<String>,
-    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<i32>]) -> Self {
-        // exterior
-        let v = points
-            .iter()
-            .map(|p| coord! { x: p.x as f64, y: p.y as f64})
-            .collect::<Vec<_>>();
-        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<u8> {
+        // 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<Bbox> {
-        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<Mbr> {
-        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
-    }
 }

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<Coord> {
         self.ls.0.clone()
     }

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;

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<String>,
+    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<i32>]) -> Self {
+        // exterior
+        let v = points
+            .iter()
+            .map(|p| coord! { x: p.x as f64, y: p.y as f64})
+            .collect::<Vec<_>>();
+        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<String>) -> 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<Bbox> {
+        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<Mbr> {
+        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
+    }
+}

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<Vec<Bbox>>,
     keypoints: Option<Vec<Vec<Keypoint>>>,
     mbrs: Option<Vec<Mbr>>,
-    masks: Option<Vec<Mask>>,
+    polygons: Option<Vec<Polygon>>,
     texts: Option<Vec<String>>,
-    pixels: Option<Vec<u8>>,
+    masks: Option<Vec<Mask>>,
 }
 
 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<u8>> {
-        self.pixels.as_ref()
+    pub fn masks(&self) -> Option<&Vec<Mask>> {
+        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<Mask>> {
-        self.masks.as_ref()
+    pub fn polygons(&self) -> Option<&Vec<Polygon>> {
+        self.polygons.as_ref()
     }
 
     pub fn bboxes(&self) -> Option<&Vec<Bbox>> {