diff --git a/examples/db/README.md b/examples/db/README.md index 9e19375..0fde408 100644 --- a/examples/db/README.md +++ b/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) diff --git a/examples/db/main.rs b/examples/db/main.rs index 13bdb87..8f797e7 100644 --- a/examples/db/main.rs +++ b/examples/db/main.rs @@ -4,9 +4,33 @@ use usls::{models::DB, Annotator, DataLoader, Options}; #[derive(argh::FromArgs)] /// Example struct Args { + /// model file + #[argh(option)] + model: Option, + /// 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<()> { @@ -14,23 +38,26 @@ fn main() -> Result<()> { .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) .with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339()) .init(); - let args: Args = argh::from_env(); // build model - 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 @@ -38,12 +65,19 @@ fn main() -> Result<()> { // 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(()) } diff --git a/examples/fast/main.rs b/examples/fast/main.rs index 84872d1..08c8001 100644 --- a/examples/fast/main.rs +++ b/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 diff --git a/examples/slanet/main.rs b/examples/slanet/main.rs index 9707c53..fe4cef1 100644 --- a/examples/slanet/main.rs +++ b/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() diff --git a/examples/svtr/main.rs b/examples/svtr/main.rs index 18704f8..12619e7 100644 --- a/examples/svtr/main.rs +++ b/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(()) diff --git a/src/models/db/config.rs b/src/models/db/config.rs index 9d5b960..c594526 100644 --- a/src/models/db/config.rs +++ b/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 { diff --git a/src/models/db/impl.rs b/src/models/db/impl.rs index de1e793..72492eb 100644 --- a/src/models/db/impl.rs +++ b/src/models/db/impl.rs @@ -2,8 +2,9 @@ use aksr::Builder; use anyhow::Result; use image::DynamicImage; use ndarray::Axis; +use rayon::prelude::*; -use crate::{elapsed, DynConf, Engine, Mbr, Ops, Options, Polygon, Processor, Ts, Xs, Ys, Y}; +use crate::{elapsed, Bbox, DynConf, Engine, Mbr, Ops, Options, Polygon, Processor, Ts, Xs, Ys, Y}; #[derive(Debug, Builder)] pub struct DB { @@ -73,102 +74,120 @@ impl DB { Ok(ys) } - pub fn summary(&mut self) { - self.ts.summary(); - } - pub fn postprocess(&mut self, xs: Xs) -> Result { - 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_polygons: Vec = Vec::new(); - let mut y_mbrs: Vec = Vec::new(); - - // input image - let (image_height, image_width) = self.processor.image0s_size[idx]; - - // reshape - let ratio = self.processor.scale_factors_hw[idx][0]; - let v = luma - .into_owned() - .into_raw_vec_and_offset() - .0 - .iter() - .map(|x| { - if x <= &self.binary_thresh { - 0u8 - } else { - (*x * 255.0) as u8 - } - }) - .collect::>(); - - let luma = Ops::resize_luma8_u8( - &v, - self.width as _, - self.height as _, - image_width as _, - image_height as _, - true, - "Bilinear", - )?; - let mask_im: image::ImageBuffer, Vec<_>> = - match image::ImageBuffer::from_raw(image_width as _, image_height as _, luma) { - None => continue, - Some(x) => x, - }; - - // contours - let contours: Vec> = - imageproc::contours::find_contours_with_threshold(&mask_im, 1); - - // loop - for contour in contours.iter() { - if contour.border_type == imageproc::contours::BorderType::Hole - && contour.points.len() <= 2 - { - continue; - } - - let polygon = Polygon::default().with_points_imageproc(&contour.points); - let delta = polygon.area() * ratio.round() as f64 * self.unclip_ratio as f64 - / polygon.perimeter(); - - // TODO: optimize - let polygon = polygon - .unclip(delta, image_width as f64, image_height as f64) - .resample(50) - // .simplify(6e-4) - .convex_hull() - .verify(); - - if let Some(bbox) = polygon.bbox() { - if bbox.height() < self.min_height || bbox.width() < self.min_width { - continue; - } - let confidence = polygon.area() as f32 / bbox.area(); - if confidence < self.confs[0] { - continue; - } - y_bbox.push(bbox.with_confidence(confidence).with_id(0)); - - if let Some(mbr) = polygon.mbr() { - y_mbrs.push(mbr.with_confidence(confidence).with_id(0)); - } - y_polygons.push(polygon.with_id(0)); - } else { - continue; - } - } - - ys.push( - Y::default() - .with_bboxes(&y_bbox) - .with_polygons(&y_polygons) - .with_mbrs(&y_mbrs), - ); - } + let ys: Vec = xs[0] + .axis_iter(Axis(0)) + .into_par_iter() + .enumerate() + .filter_map(|(idx, luma)| { + // input image + let (image_height, image_width) = self.processor.image0s_size[idx]; + + // reshape + let ratio = self.processor.scale_factors_hw[idx][0]; + let v = luma + .as_slice()? + .par_iter() + .map(|x| { + if x <= &self.binary_thresh { + 0u8 + } else { + (*x * 255.0) as u8 + } + }) + .collect::>(); + + let luma = Ops::resize_luma8_u8( + &v, + self.width as _, + self.height as _, + image_width as _, + image_height as _, + true, + "Bilinear", + ) + .ok()?; + let mask_im: image::ImageBuffer, Vec<_>> = + image::ImageBuffer::from_raw(image_width as _, image_height as _, luma)?; + + // contours + let contours: Vec> = + imageproc::contours::find_contours_with_threshold(&mask_im, 1); + + // loop + let (y_polygons, y_bbox, y_mbrs): (Vec, Vec, Vec) = contours + .par_iter() + .filter_map(|contour| { + if contour.border_type == imageproc::contours::BorderType::Hole + && contour.points.len() <= 2 + { + return None; + } + + let polygon = Polygon::default() + .with_points_imageproc(&contour.points) + .with_id(0); + let delta = + polygon.area() * ratio.round() as f64 * self.unclip_ratio as f64 + / polygon.perimeter(); + + let polygon = polygon + .unclip(delta, image_width as f64, image_height as f64) + .resample(50) + // .simplify(6e-4) + .convex_hull() + .verify(); + + polygon.bbox().and_then(|bbox| { + if bbox.height() < self.min_height || bbox.width() < self.min_width { + return None; + } + let confidence = polygon.area() as f32 / bbox.area(); + if confidence < self.confs[0] { + return None; + } + let bbox = bbox.with_confidence(confidence).with_id(0); + let mbr = polygon + .mbr() + .map(|mbr| mbr.with_confidence(confidence).with_id(0)); + + Some((polygon, bbox, mbr)) + }) + }) + .fold( + || (Vec::new(), Vec::new(), Vec::new()), + |mut acc, (polygon, bbox, mbr)| { + acc.0.push(polygon); + acc.1.push(bbox); + if let Some(mbr) = mbr { + acc.2.push(mbr); + } + acc + }, + ) + .reduce( + || (Vec::new(), Vec::new(), Vec::new()), + |mut acc, (polygons, bboxes, mbrs)| { + acc.0.extend(polygons); + acc.1.extend(bboxes); + acc.2.extend(mbrs); + acc + }, + ); + + Some( + Y::default() + .with_bboxes(&y_bbox) + .with_polygons(&y_polygons) + .with_mbrs(&y_mbrs), + ) + }) + .collect(); Ok(ys.into()) } + + pub fn summary(&mut self) { + self.ts.summary(); + } } diff --git a/src/models/svtr/config.rs b/src/models/svtr/config.rs index 93fc38e..3d82fa3 100644 --- a/src/models/svtr/config.rs +++ b/src/models/svtr/config.rs @@ -4,6 +4,7 @@ impl crate::Options { Self::default() .with_model_name("svtr") .with_model_ixx(0, 0, (1, 1, 8).into()) + .with_model_ixx(0, 1, 3.into()) .with_model_ixx(0, 2, 48.into()) .with_model_ixx(0, 3, (320, 960, 1600).into()) .with_resize_mode(crate::ResizeMode::FitHeight)