Using rayon to accelerate DB post-processing

examples/db/README.md

@@ -1,10 +1,13 @@
 ## Quick Start
 
 ```shell
-cargo run -r --example db
+cargo run -r -F cuda --example db -- --device cuda --dtype fp16
 ```
 
 ## Results
 
 ![](https://github.com/jamjamjon/assets/releases/download/db/demo-paper.png)
+![](https://github.com/jamjamjon/assets/releases/download/db/demo-slanted-text-number.png)
+![](https://github.com/jamjamjon/assets/releases/download/db/demo-table-en.png)
+![](https://github.com/jamjamjon/assets/releases/download/db/demo-table-ch.png)
 ![](https://github.com/jamjamjon/assets/releases/download/db/demo-sign.png)

examples/db/main.rs

@@ -4,46 +4,80 @@ use usls::{models::DB, Annotator, DataLoader, Options};
 #[derive(argh::FromArgs)]
 /// Example
 struct Args {
+    /// model file
+    #[argh(option)]
+    model: Option<String>,
+
     /// device
     #[argh(option, default = "String::from(\"cpu:0\")")]
     device: String,
+
+    /// dtype
+    #[argh(option, default = "String::from(\"auto\")")]
+    dtype: String,
+
+    /// show bboxes
+    #[argh(option, default = "false")]
+    show_bboxes: bool,
+
+    /// show mbrs
+    #[argh(option, default = "false")]
+    show_mbrs: bool,
+
+    /// show bboxes confidence
+    #[argh(option, default = "false")]
+    show_bboxes_conf: bool,
+
+    /// show mbrs confidence
+    #[argh(option, default = "false")]
+    show_mbrs_conf: bool,
 }
 
 fn main() -> Result<()> {
     tracing_subscriber::fmt()
         .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
         .with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
         .init();
-
     let args: Args = argh::from_env();
 
     // build model
-    let options = Options::ppocr_det_v4_server_ch()
-        .with_model_device(args.device.as_str().try_into()?)
-        .commit()?;
-    let mut model = DB::new(options)?;
+    let options = match &args.model {
+        Some(m) => Options::db().with_model_file(m),
+        None => Options::ppocr_det_v4_ch().with_model_dtype(args.dtype.as_str().try_into()?),
+    };
+    let mut model = DB::new(
+        options
+            .with_model_device(args.device.as_str().try_into()?)
+            .commit()?,
+    )?;
 
     // load image
     let x = DataLoader::try_read_batch(&[
+        "images/db.png",
         "images/table.png",
-        "images/table1.jpg",
-        "images/table2.png",
         "images/table-ch.jpg",
-        "images/db.png",
         "images/street.jpg",
+        "images/slanted-text-number.jpg",
     ])?;
 
     // run
     let y = model.forward(&x)?;
 
     // annotate
     let annotator = Annotator::default()
-        .without_bboxes(true)
-        .without_mbrs(true)
+        .without_bboxes(!args.show_bboxes)
+        .without_mbrs(!args.show_mbrs)
+        .without_bboxes_name(true)
+        .without_mbrs_name(true)
+        .without_bboxes_conf(!args.show_bboxes_conf)
+        .without_mbrs_conf(!args.show_mbrs_conf)
         .with_polygons_alpha(60)
         .with_contours_color([255, 105, 180, 255])
         .with_saveout(model.spec());
     annotator.annotate(&x, &y);
 
+    // summary
+    model.summary();
+
     Ok(())
 }

examples/fast/main.rs

@@ -41,12 +41,11 @@ fn main() -> Result<()> {
 
     // load image
     let x = DataLoader::try_read_batch(&[
+        "images/db.png",
         "images/table.png",
-        "images/table1.jpg",
-        "images/table2.png",
         "images/table-ch.jpg",
-        "images/db.png",
         "images/street.jpg",
+        "images/slanted-text-number.jpg",
     ])?;
 
     // run

examples/slanet/main.rs

@@ -11,6 +11,10 @@ struct Args {
     /// device
     #[argh(option, default = "String::from(\"cpu:0\")")]
     device: String,
+
+    /// dtype
+    #[argh(option, default = "String::from(\"auto\")")]
+    dtype: String,
 }
 
 fn main() -> Result<()> {
@@ -24,6 +28,7 @@ fn main() -> Result<()> {
     // build model
     let options = Options::slanet_lcnet_v2_mobile_ch()
         .with_model_device(args.device.as_str().try_into()?)
+        .with_model_dtype(args.dtype.as_str().try_into()?)
         .commit()?;
     let mut model = SLANet::new(options)?;
 
@@ -32,7 +37,7 @@ fn main() -> Result<()> {
 
     // run
     let ys = model.forward(&xs)?;
-    println!("{:?}", ys);
+    // println!("{:?}", ys);
 
     // annotate
     let annotator = Annotator::default()

examples/svtr/main.rs

@@ -7,6 +7,10 @@ struct Args {
     /// device
     #[argh(option, default = "String::from(\"cpu:0\")")]
     device: String,
+
+    /// dtype
+    #[argh(option, default = "String::from(\"auto\")")]
+    dtype: String,
 }
 
 fn main() -> Result<()> {
@@ -19,9 +23,9 @@ fn main() -> Result<()> {
 
     // build model
     let options = Options::ppocr_rec_v4_ch()
-        // svtr_v2_teacher_ch()
-        // .with_batch_size(2)
+        // repsvtr_ch()
         .with_model_device(args.device.as_str().try_into()?)
+        .with_model_dtype(args.dtype.as_str().try_into()?)
         .commit()?;
     let mut model = SVTR::new(options)?;
 
@@ -37,7 +41,7 @@ fn main() -> Result<()> {
         println!("{paths:?}: {:?}", ys)
     }
 
-    //summary
+    // summary
     model.summary();
 
     Ok(())

src/models/db/config.rs

@@ -33,8 +33,6 @@ impl crate::Options {
         Self::db()
             .with_image_mean(&[0.798, 0.785, 0.772])
             .with_image_std(&[0.264, 0.2749, 0.287])
-        // .with_binary_thresh(0.3)
-        // .with_class_confs(&[0.1])
     }
 
     pub fn db_mobilenet_v3_large() -> Self {